diff --git a/dynamax/generalized_gaussian_ssm/inference_test.py b/dynamax/generalized_gaussian_ssm/inference_test.py index 0aeaeec69..47f2358aa 100644 --- a/dynamax/generalized_gaussian_ssm/inference_test.py +++ b/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) diff --git a/dynamax/linear_gaussian_ssm/inference_test.py b/dynamax/linear_gaussian_ssm/inference_test.py index a1c89e06c..07aef0e8f 100644 --- a/dynamax/linear_gaussian_ssm/inference_test.py +++ b/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) diff --git a/dynamax/linear_gaussian_ssm/info_inference_test.py b/dynamax/linear_gaussian_ssm/info_inference_test.py index 2c8175785..d9ab8672c 100644 --- a/dynamax/linear_gaussian_ssm/info_inference_test.py +++ b/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) diff --git a/dynamax/nonlinear_gaussian_ssm/inference_ekf_test.py b/dynamax/nonlinear_gaussian_ssm/inference_ekf_test.py index 8904dfe30..3fbf8fccf 100644 --- a/dynamax/nonlinear_gaussian_ssm/inference_ekf_test.py +++ b/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) diff --git a/dynamax/nonlinear_gaussian_ssm/inference_ukf_test.py b/dynamax/nonlinear_gaussian_ssm/inference_ukf_test.py index cb0c649ab..43060171c 100644 --- a/dynamax/nonlinear_gaussian_ssm/inference_ukf_test.py +++ b/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) diff --git a/dynamax/parameters_test.py b/dynamax/parameters_test.py index 5139a3361..1b953b423 100644 --- a/dynamax/parameters_test.py +++ b/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(): diff --git a/dynamax/utils/distributions_test.py b/dynamax/utils/distributions_test.py index 552c9e63e..556ffdabf 100644 --- a/dynamax/utils/distributions_test.py +++ b/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)) diff --git a/dynamax/utils/utils.py b/dynamax/utils/utils.py index 3dd439c73..afed880cf 100644 --- a/dynamax/utils/utils.py +++ b/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