diff --git a/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py b/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py index 1621a73a1..fbc55bd28 100644 --- a/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py +++ b/test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py @@ -8,12 +8,12 @@ _times_from_positions, ) -# `_times_from_positions` is jax-backed; it needs jax installed to run (jax -# ships via the `[optional]` extras). The NumPy-only Python-version matrix has -# no jax, so skip there rather than fail. -requires_jax = pytest.mark.skipif( - importlib.util.find_spec("jax") is None, - reason="requires jax (installed via the [optional] extras; absent on the NumPy-only matrix env)", +# `_times_from_positions` runs blackjax (which pulls jax); both ship via the +# `[optional]` extras. The Python-version matrix installs autofit without those +# extras, so skip there rather than fail with `No module named 'blackjax'`. +requires_blackjax = pytest.mark.skipif( + importlib.util.find_spec("blackjax") is None, + reason="requires blackjax (installed via the [optional] extras; absent on the matrix env)", ) pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") @@ -101,7 +101,7 @@ def test__identifier_fields_distinguish_run_shape(): assert (c.num_warmup, c.num_samples, c.num_chains) == (100, 999, 1) -@requires_jax +@requires_blackjax def test__times_from_positions_clamps_low_ess(): # Construct a degenerate "chain" — every sample is identical → ESS # collapses; the clamp must keep ``times`` finite and equal to N. @@ -118,7 +118,7 @@ def test__times_from_positions_clamps_low_ess(): assert np.all(times <= n_samples + 1e-9) -@requires_jax +@requires_blackjax def test__times_from_positions_independent_chain(): # Independent draws → ESS ≈ N, so τ ≈ 1. rng = np.random.default_rng(1)