test: remove jax from conftest, gate jax test files with importorskip

Per the 'library unit tests stay numpy-only' rule, conftest must not
import jax at all — not even via try/except.

Changes:
- conftest.py: remove `import jax`, drop the recreate fixture (moved
  into test_pytrees.py since that's the only consumer), add a
  find_spec-gated collect_ignore_glob that skips test_autofit/jax/
  and test_jacobians.py when jax isn't installed.
- jax/test_pytrees.py: move pytest.importorskip("jax") BEFORE the
  jax imports (was AFTER, so module load failed on no-jax). Inline
  the recreate fixture here.
- graphical/functionality/test_jacobians.py: add pytest.importorskip
  at module top so the file is self-contained even if conftest's
  collect_ignore is bypassed.

This supersedes 1692c8e (try/except hack).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 2 people

test_autofit/conftest.py

@@ -1,7 +1,4 @@
-try:
-    import jax
-except ImportError:
-    jax = None
+import importlib.util
 import multiprocessing
 import os
 import shutil
@@ -19,23 +16,22 @@
 from autofit.database.model import sa
 from autofit.non_linear.search import abstract_search
 
+# Skip JAX-only tests when jax isn't installed. find_spec checks availability
+# WITHOUT importing jax, so this conftest stays numpy-only per the
+# "library unit tests stay numpy-only" rule.
+collect_ignore_glob = []
+if importlib.util.find_spec("jax") is None:
+    collect_ignore_glob = [
+        "jax/*.py",
+        "graphical/functionality/test_jacobians.py",
+    ]
+
 if sys.platform == "darwin":
     multiprocessing.set_start_method("fork")
 
 directory = Path(__file__).parent
 
 
-@pytest.fixture(name="recreate")
-def recreate():
-
-    def _recreate(o):
-        flatten_func, unflatten_func = jax._src.tree_util._registry[type(o)]
-        children, aux_data = flatten_func(o)
-        return unflatten_func(aux_data, children)
-
-    return _recreate
-
-
 @pytest.fixture(autouse=True)
 def turn_off_gc(monkeypatch):
     monkeypatch.setattr(abstract_search, "gc", MagicMock())

test_autofit/graphical/functionality/test_jacobians.py

@@ -1,8 +1,11 @@
 from itertools import combinations
-import jax
 import numpy as np
 import pytest
 
+# Belt-and-suspenders: the top-level conftest also ignores this file when
+# jax isn't installed; importorskip keeps the file self-contained.
+jax = pytest.importorskip("jax")
+
 from autofit.mapper.variable import variables
 from autofit.graphical.factor_graphs import (
     Factor,

test_autofit/jax/test_pytrees.py

@@ -1,20 +1,36 @@
 import numpy as np
 import pytest
-import jax.numpy as jnp
+
+# importorskip MUST run before any jax import — skips the whole module
+# cleanly when jax isn't installed (Python <3.11 in our [jax] extra gate).
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
 from jax.tree_util import register_pytree_node_class
 
 import autofit as af
 from autofit import UniformPrior
 
-jax = pytest.importorskip("jax")
-
 UniformPrior = register_pytree_node_class(UniformPrior)
 GaussianPrior = register_pytree_node_class(af.GaussianPrior)
 TruncatedGaussianPrior = register_pytree_node_class(af.TruncatedGaussianPrior)
 Collection = register_pytree_node_class(af.Collection)
 Model = register_pytree_node_class(af.Model)
 ModelInstance = register_pytree_node_class(af.ModelInstance)
 
+
+@pytest.fixture(name="recreate")
+def make_recreate():
+    """jax-pytree-roundtrip fixture, scoped to this test module since it's
+    the only consumer and it must not pollute the numpy-only top-level
+    conftest."""
+
+    def _recreate(o):
+        flatten_func, unflatten_func = jax._src.tree_util._registry[type(o)]
+        children, aux_data = flatten_func(o)
+        return unflatten_func(aux_data, children)
+
+    return _recreate
+
 @pytest.fixture(name="gaussian")
 def make_gaussian():
     return af.ex.Gaussian(centre=1.0, sigma=1.0, normalization=1.0)