fix(jax): make jax usage truly optional across library

The Python version matrix CI was failing on 3.9 and 3.10 with
`ModuleNotFoundError: No module named 'jax'` raised from autofit
library code, even though the [jax] extra is gated `python_version
>= '3.11'`. Three eager / mis-polarized jax paths in autofit/ caused
this:

* `factor.py:_set_jacobians` did `import jax` at the top of the
  function, but only the `jacfwd` / `jacobian` fallbacks need it.
  The default `numerical_jacobian=True` path that every plain
  `graph.Factor(...)` constructor hits never touches jax. Push
  `import jax` down into the two branches that actually use it.

* `messages/normal.py:value_for` and `mapper/prior_model/array.py`
  both routed plain Python `float` / `int` to `import jax.numpy`
  (their `else` branch). After this fix the isinstance checks include
  `(np.ndarray, np.float64, float, int, list)` so only genuine jax
  values fall through to the jax branch.

Five unit tests in `test_autofit/` exercised jax-only behaviour
(`Fitness(use_jax_jit=True)`, `_JitFittableAnalysis` JIT dispatch,
`af.Array` returning a `jnp.ndarray`). They are moved out of the
unit suite to `autofit_workspace_test/scripts/graphical/jax_assertions.py`
in line with the policy that all library code paths must support
pure numpy.

Verified locally:
* `import autofit` succeeds in a fresh venv with no jax installed.
* `pytest test_autofit/` — 1225 passed, 1 skipped, 0 failed without
  jax; 1245 passed, 0 failed with jax.

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

Author: 2 people

autofit/graphical/factor_graphs/factor.py

@@ -284,8 +284,6 @@ def _set_jacobians(
         numerical_jacobian=True,
         jacfwd=True,
     ):
-        import jax
-
         self._vjp = vjp
         self._jacfwd = jacfwd
         if vjp or factor_vjp:
@@ -302,8 +300,10 @@ def _set_jacobians(
             elif numerical_jacobian:
                 self._factor_jacobian = self._numerical_factor_jacobian
             elif jacfwd:
+                import jax
                 self._jacobian = jax.jacfwd(self._factor, range(self.n_args))
             else:
+                import jax
                 self._jacobian = jax.jacobian(self._factor, range(self.n_args))
 
     def _factor_value(self, raw_fval) -> FactorValue:

autofit/mapper/prior_model/array.py

@@ -87,15 +87,15 @@ def _instance_for_arguments(
                 pass
 
             if make_array:
-                if isinstance(value, np.ndarray) or isinstance(value, np.float64):
+                if isinstance(value, (np.ndarray, np.float64, float, int)):
                     array = np.zeros(self.shape)
                     make_array = False
                 else:
                     import jax.numpy as jnp
                     array = jnp.zeros(self.shape)
                     make_array = False
 
-            if isinstance(value, np.ndarray) or isinstance(value, np.float64):
+            if isinstance(value, (np.ndarray, np.float64, float, int)):
                 array[index] = value
             else:
                 array = array.at[index].set(value)

autofit/messages/normal.py

@@ -417,7 +417,7 @@ def value_for(self, unit: float) -> float:
         >>> prior = af.GaussianPrior(mean=1.0, sigma=2.0)
         >>> physical_value = prior.value_for(unit=0.5)
         """
-        if isinstance(unit, np.ndarray) or isinstance(unit, np.float64):
+        if isinstance(unit, (np.ndarray, np.float64, float, int, list)):
             from scipy.special import erfinv as scipy_erfinv
             inv = scipy_erfinv(1 - 2.0 * (1.0 - unit))
         else:

test_autofit/analysis/test_use_jax_for_visualization.py

@@ -20,28 +20,6 @@ def fit_from(self, instance):
         return ("fit", instance)
 
 
-class _JitFittableAnalysis(af.Analysis):
-    """Analysis with a ``fit_from`` returning a JIT-traceable array.
-
-    ``_FittableAnalysis.fit_from`` returns a Python tuple with a string literal,
-    which is not tracer-compatible. For the JIT-enabled dispatch path we need a
-    ``fit_from`` whose output is entirely JAX-compatible.
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.fit_from_calls = 0
-
-    def log_likelihood_function(self, instance):
-        return 0.0
-
-    def fit_from(self, instance):
-        import jax.numpy as jnp
-
-        self.fit_from_calls += 1
-        return jnp.asarray(instance) * 2.0
-
-
 def test_default_flag_is_false():
     analysis = af.Analysis()
     assert analysis._use_jax is False
@@ -70,23 +48,6 @@ def test_pyauto_disable_jax_env_var_clears_both_flags(monkeypatch):
     assert analysis._use_jax_for_visualization is False
 
 
-def test_fit_for_visualization_dispatches_through_jit_when_flag_set():
-    import jax.numpy as jnp
-
-    analysis = _JitFittableAnalysis(use_jax=True, use_jax_for_visualization=True)
-
-    assert getattr(analysis, "_jitted_fit_from", None) is None
-
-    result_1 = analysis.fit_for_visualization(instance=1.0)
-    assert analysis._jitted_fit_from is not None
-    assert jnp.allclose(result_1, jnp.asarray(2.0))
-
-    jitted_after_first = analysis._jitted_fit_from
-    result_2 = analysis.fit_for_visualization(instance=3.0)
-    assert analysis._jitted_fit_from is jitted_after_first
-    assert jnp.allclose(result_2, jnp.asarray(6.0))
-
-
 def test_fit_for_visualization_works_without_flag():
     analysis = _FittableAnalysis()
     result = analysis.fit_for_visualization(instance="sentinel")

test_autofit/mapper/test_array.py

@@ -159,41 +159,3 @@ def test_tree_flatten(array):
     ).all()
 
 
-class Analysis(af.Analysis):
-    def log_likelihood_function(self, instance):
-        return -float(
-            np.mean(
-                (
-                    np.array(
-                        [
-                            [0.1, 0.2],
-                            [0.3, 0.4],
-                        ]
-                    )
-                    - instance
-                )
-                ** 2
-            )
-        )
-
-
-def test_optimisation():
-
-    import jax.numpy as jnp
-
-    array = af.Array(
-        shape=(2, 2),
-        prior=af.UniformPrior(
-            lower_limit=0.0,
-            upper_limit=1.0,
-        ),
-    )
-    result = af.DynestyStatic().fit(model=array, analysis=Analysis())
-
-    posterior = result.model
-    array[0, 0] = posterior[0, 0]
-    array[0, 1] = posterior[0, 1]
-
-    result = af.DynestyStatic().fit(model=array, analysis=Analysis())
-
-    assert isinstance(result.instance, jnp.ndarray)

test_autofit/non_linear/test_fitness_jax_dispatch.py

@@ -23,39 +23,6 @@ def test_default_dispatch_is_call():
     assert fitness.use_jax_vmap is False
 
 
-def test_jit_dispatch_sets_call_to_jit():
-    fitness = _make_fitness(use_jax_jit=True)
-    assert fitness.use_jax_jit is True
-    assert fitness._call is fitness._jit
-
-
-def test_vmap_takes_precedence_over_jit():
-    fitness = _make_fitness(use_jax_jit=True, use_jax_vmap=True)
-    assert fitness._call is fitness._vmap
-
-
-def test_pickle_strips_jax_cached_attrs():
-    """
-    Dynesty's checkpoint writes pickle the loglikelihood. JAX-compiled
-    callables (jax.jit / jax.vmap / jax.grad) carry C++ XLA state that
-    cannot roundtrip through pickle. ``Fitness.__getstate__`` must drop
-    them; ``Fitness.__setstate__`` re-derives the dispatch on resume.
-    """
-    fitness = _make_fitness(use_jax_jit=True)
-
-    state = fitness.__getstate__()
-    assert "_call" not in state
-    assert "_jit" not in state
-    assert "_vmap" not in state
-    assert "_grad" not in state
-
-    blob = pickle.dumps(fitness)
-    restored = pickle.loads(blob)
-
-    assert restored.use_jax_jit is True
-    assert restored._call is restored._jit
-
-
 def test_pickle_default_path_unchanged():
     fitness = _make_fitness()