Fix TPU tests

dynamax/generalized_gaussian_ssm/inference_test.py

@@ -12,7 +12,7 @@
 from functools import partial
 
 if has_tpu():
-    allclose = partial(jnp.allclose, atol=1e-1)
+    allclose = partial(jnp.allclose, atol=1e-3)
 else:
     allclose = partial(jnp.allclose, atol=1e-3)
 

dynamax/linear_gaussian_ssm/inference_test.py

@@ -15,7 +15,7 @@
 
 # Use different tolerance threshold for TPU
 if has_tpu():
-    allclose = partial(jnp.allclose, atol=1e-1)
+    allclose = partial(jnp.allclose, atol=1e-3)
 else:
     allclose = partial(jnp.allclose, atol=1e-4)
 

dynamax/linear_gaussian_ssm/info_inference_test.py

@@ -14,7 +14,7 @@
 
 # Use lower tolerance for TPU tests.
 if has_tpu():
-    allclose = partial(jnp.allclose, atol=1e-1)
+    allclose = partial(jnp.allclose, atol=1e-3)
 else:
     allclose = partial(jnp.allclose, atol=1e-4)
 

dynamax/nonlinear_gaussian_ssm/inference_ekf_test.py

@@ -14,7 +14,7 @@
 
 if has_tpu():
     # TPU has very poor numerical stability
-    allclose = partial(jnp.allclose, atol=1e-1)
+    allclose = partial(jnp.allclose, atol=1e-3)
 else:
     allclose = partial(jnp.allclose, atol=1e-4)
 

dynamax/nonlinear_gaussian_ssm/inference_ukf_test.py

@@ -10,7 +10,7 @@
 from functools import partial
 
 if has_tpu():
-    allclose = partial(jnp.allclose, atol=1e-1)
+    allclose = partial(jnp.allclose, atol=1e-3)
 else:
     allclose = partial(jnp.allclose, atol=1e-4)
 

dynamax/parameters_test.py

@@ -6,6 +6,7 @@
 import tensorflow_probability.substrates.jax.bijectors as tfb
 
 from dynamax.parameters import ParameterProperties, to_unconstrained, from_unconstrained, log_det_jac_constrain
+from dynamax.utils.utils import has_tpu
 from dynamax.parameters import trainable_flags, ensure_all_or_none_trainable
 from jax import jit, value_and_grad, lax
 from jax.tree_util import tree_map, tree_leaves
@@ -149,7 +150,8 @@ def step(carry, args):
     assert jnp.allclose(params.initial.probs, original_params.initial.probs)
     assert not jnp.allclose(params.transitions.transition_matrix, original_params.transitions.transition_matrix)
     assert not jnp.allclose(params.emissions.means, original_params.emissions.means)
-    assert jnp.allclose(params.emissions.scales, original_params.emissions.scales)
+    assert jnp.allclose(params.emissions.scales, original_params.emissions.scales,
+                        atol=1e-3 if has_tpu() else 1e-8)
 
 
 def test_logdet_jacobian():

dynamax/utils/distributions_test.py

@@ -15,6 +15,7 @@
 from dynamax.utils.distributions import MatrixNormalPrecision
 from dynamax.utils.distributions import NormalInverseGamma
 from dynamax.utils.distributions import NormalInverseWishart
+from dynamax.utils.utils import has_tpu
 
 tfd = tfp.distributions
 tfb = tfp.bijectors
@@ -71,8 +72,10 @@ def test_inverse_wishart_variance_vectorization():
        [-0.8869951,  5.07704  , -0.8494578],
        [-0.9199044, -0.8494578,  4.1288185]], dtype=jnp.float32)
     𝜈2 = 8.0  # >p + 3
-    assert all(jnp.linalg.eigvals(Ψ1) > 0)
-    assert all(jnp.linalg.eigvals(Ψ2) > 0)
+    assert jnp.allclose(Ψ1, Ψ1.T)  # eigvalsh assumes symmetry.
+    assert jnp.allclose(Ψ2, Ψ2.T)
+    assert all(jnp.linalg.eigvalsh(Ψ1) > 0)
+    assert all(jnp.linalg.eigvalsh(Ψ2) > 0)
 
     # Make a (2, 1) batch shape to test vectorization over >1 leading axis.
     𝜈 = jnp.stack([𝜈1, 𝜈2]) # Shape: (2,)
@@ -95,7 +98,8 @@ def test_inverse_wishart_sample_non_diagonal_scale(n_samples: int = 10_000, num_
     Ψ = jnp.array([[20.712932, 25.124634],
         [25.124634, 32.814785]], dtype=jnp.float32)  # k x k
     Ψ_diag = jnp.diagonal(Ψ)
-    assert all(jnp.linalg.eigvals(Ψ) > 0)  # Is positive definite.
+    assert jnp.allclose(Ψ, Ψ.T)  # eigvalsh assumes symmetry.
+    assert all(jnp.linalg.eigvalsh(Ψ) > 0)  # Is positive definite.
 
     iw = InverseWishart(df=𝜈, scale=Ψ)
     Σs = iw.sample(sample_shape=n_samples, seed=jr.key(42))
@@ -187,7 +191,7 @@ def test_matrix_normal_inverse_wishart_log_prob(
     lp_mn = jnp.array([matrix_normal.logpdf(m, loc, sigma, jnp.linalg.inv(col_precision)) \
                        for m, sigma in zip(Matrix_samples, Sigma_samples)])
 
-    assert jnp.allclose(mniw_log_probs, lp_iw + lp_mn)
+    assert jnp.allclose(mniw_log_probs, lp_iw + lp_mn, atol=1e-3 if has_tpu() else 1e-8)
     assert jnp.allclose(mniw_probs, jnp.exp(lp_iw + lp_mn))
 
 

dynamax/utils/utils.py

@@ -2,7 +2,6 @@
 Utility functions for the library.
 """
 import jax
-import jaxlib
 import jax.numpy as jnp
 import jax.random as jr
 
@@ -16,10 +15,7 @@
 
 def has_tpu():
     """Check if the current device is a TPU."""
-    try:
-        return isinstance(jax.devices()[0], jaxlib.xla_extension.TpuDevice)
-    except:
-        return False
+    return jax.default_backend() == "tpu"
 
 
 @jit