fix(inversion): review findings (#566)

Move the no-module-level-jax import guard into its own test module: the
module-level skip in test_jax_active_set.py (applied when jax is absent)
also skipped the guard, contradicting its comment, so the guard never ran
on the NumPy-only env it exists for.

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

Author: Jammy2211

test_autoarray/util/test_jax_active_set.py

@@ -15,33 +15,13 @@
 """
 
 import importlib
-import sys
 
 import numpy as np
 import pytest
 
 
-def test__jax_active_set_module_never_imports_jax_at_module_level(monkeypatch):
-    # Library unit tests are NumPy-only: importing the module must succeed even when jax is
-    # unimportable. A None entry in sys.modules makes any `import jax` raise ImportError, so a
-    # module-level import would fail this reload.
-    monkeypatch.setitem(sys.modules, "jax", None)
-    monkeypatch.setitem(sys.modules, "jaxnnls", None)
-
-    module = importlib.reload(importlib.import_module("autoarray.util.jax_active_set"))
-
-    for name in (
-        "masked_solve",
-        "certify",
-        "active_set_search",
-        "solve_certified",
-        "solve_certified_with_fallback",
-    ):
-        assert hasattr(module, name)
-
-
-# jax is an `[optional]` extra and is absent on the NumPy-only matrix env: every test below the import test
-# skips there, while the import test above still runs.
+# jax is an `[optional]` extra and is absent on the NumPy-only matrix env: every test in this module skips
+# there. The no-module-level-jax-import guard lives in `test_jax_active_set_import.py` so it still runs.
 if importlib.util.find_spec("jax") is None:
     pytestmark = pytest.mark.skip(reason="requires jax (the [optional] extras)")
 

test_autoarray/util/test_jax_active_set_import.py

@@ -0,0 +1,26 @@
+"""
+Import guard for `autoarray.util.jax_active_set`, kept apart from `test_jax_active_set.py` because that module
+skips wholesale when jax is absent, and this test must run exactly there.
+"""
+
+import importlib
+import sys
+
+
+def test__jax_active_set_module_never_imports_jax_at_module_level(monkeypatch):
+    # Library unit tests are NumPy-only: importing the module must succeed even when jax is
+    # unimportable. A None entry in sys.modules makes any `import jax` raise ImportError, so a
+    # module-level import would fail this reload.
+    monkeypatch.setitem(sys.modules, "jax", None)
+    monkeypatch.setitem(sys.modules, "jaxnnls", None)
+
+    module = importlib.reload(importlib.import_module("autoarray.util.jax_active_set"))
+
+    for name in (
+        "masked_solve",
+        "certify",
+        "active_set_search",
+        "solve_certified",
+        "solve_certified_with_fallback",
+    ):
+        assert hasattr(module, name)