Use precision=HIGHEST in expm repeated squaring (#3601)
Use precision=HIGHEST in expm repeated squaring (#3601)

* Use precision=HIGHEST in expm repeated squaring

This will improve the output accuracy on TPUs, and doesn't currently affect other platforms.

Also remove a spurious duplicate line.

* Also use HIGHEST precision in Pade approximants

* Update linalg.py

* Update linalg.py
tensorflow/probability,google/jax,google/jax,google/jax,google/jax,tensorflow/probability