test: skip jax-backed unit tests when jax is absent (Python matrix)

test_autofit/analysis/test_latent_variables.py

@@ -1,9 +1,19 @@
+import importlib.util
+
 import numpy as np
 import pytest
 
 import autofit as af
 from autoconf.conf import with_config
 from autofit import DirectoryPaths, SamplesPDF
+
+# This test drives the `use_jax=True` path, which needs jax installed (it 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)",
+)
 from autofit.text.text_util import result_info_from
 
 
@@ -104,6 +114,7 @@ def test_latent_batch_mode_default_is_vmap():
     assert Analysis.LATENT_BATCH_MODE == "vmap"
 
 
+@requires_jax
 def test_latent_batch_mode_invalid_value_raises_clear_error():
     """
     Misspelt `LATENT_BATCH_MODE` values should fail with a clear ValueError

test_autofit/non_linear/search/mcmc/test_blackjax_nuts.py

@@ -1,3 +1,4 @@
+import importlib.util
 import os
 import numpy as np
 import pytest
@@ -7,6 +8,14 @@
     _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)",
+)
+
 pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning")
 
 
@@ -92,6 +101,7 @@ def test__identifier_fields_distinguish_run_shape():
     assert (c.num_warmup, c.num_samples, c.num_chains) == (100, 999, 1)
 
 
+@requires_jax
 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.
@@ -108,6 +118,7 @@ def test__times_from_positions_clamps_low_ess():
     assert np.all(times <= n_samples + 1e-9)
 
 
+@requires_jax
 def test__times_from_positions_independent_chain():
     # Independent draws → ESS ≈ N, so τ ≈ 1.
     rng = np.random.default_rng(1)