From ec30492b07ebf6b29714cc996de75e58c17846cd Mon Sep 17 00:00:00 2001 From: nkbranigan <81596699+nkbranigan@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:41:35 -0700 Subject: [PATCH] Fix TPU tests - Fix `has_tpu`: replace removed `jaxlib` internals with modern API - Update TPU-specific tolerances to 1e-3 - Change `eigvals`, which is not implemented on TPU, to `eigvalsh` Resolves #260, resolves #261, resolves #262 --- dynamax/generalized_gaussian_ssm/inference_test.py | 2 +- dynamax/linear_gaussian_ssm/inference_test.py | 2 +- dynamax/linear_gaussian_ssm/info_inference_test.py | 2 +- dynamax/nonlinear_gaussian_ssm/inference_ekf_test.py | 2 +- dynamax/nonlinear_gaussian_ssm/inference_ukf_test.py | 2 +- dynamax/parameters_test.py | 4 +++- dynamax/utils/distributions_test.py | 12 ++++++++---- dynamax/utils/utils.py | 6 +----- 8 files changed, 17 insertions(+), 15 deletions(-) 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 9477f10da..c6306f66b 100644 --- a/dynamax/parameters_test.py +++ b/dynamax/parameters_test.py @@ -5,6 +5,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 jax import jit, value_and_grad, lax from jax.tree_util import tree_map, tree_leaves from jaxtyping import Float, Array @@ -132,7 +133,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