feat: add autofit.jax opt-in pytree registration API

Introduces autofit.jax.enable_pytrees() and autofit.jax.register_model(model)
so callers can register PyAutoFit prior/model classes (and the user classes
referenced by model.cls) with jax.tree_util without paying an eager JAX
import at every "import autofit".

register_model walks a Model / Collection and builds per-class flatten /
unflatten functions that partition constructor arguments: constants (from
model.direct_instance_tuples) go into JAX aux_data so they stay as concrete
Python values inside jit traces, while prior-derived arguments become
children (and therefore tracers). This keeps constants like Galaxy.redshift
usable in Python control flow (e.g. sorted(...) by redshift) under JIT.

Both functions are no-ops if JAX is not installed. Re-registration of a
class that is already registered is tolerated silently.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: 2 people

autofit/jax/__init__.py

@@ -0,0 +1,16 @@
+"""Opt-in JAX integration helpers for PyAutoFit.
+
+The library defines pytree ``tree_flatten`` / ``tree_unflatten`` methods on
+its model and prior classes but, by design, does NOT register them with JAX
+at import time. Eager registration would force ``jax.tree_util`` to load on
+every ``import autofit`` and reintroduce the heavy JAX import that the
+2025-11 cleanup removed.
+
+Call :func:`enable_pytrees` once before crossing a ``jax.jit`` or
+``jax.vmap`` boundary with PyAutoFit objects. The function is a no-op if JAX
+is not installed.
+"""
+
+from .pytrees import enable_pytrees, register_model
+
+__all__ = ["enable_pytrees", "register_model"]

autofit/jax/pytrees.py

@@ -0,0 +1,150 @@
+"""Lazy JAX pytree registration for PyAutoFit's model and prior classes.
+
+The classes themselves already carry the necessary ``tree_flatten`` /
+``tree_unflatten`` methods. This module simply registers them with
+``jax.tree_util`` on demand, via the lazy autoconf wrapper, so callers can
+pass ``Model`` / ``Collection`` / ``ModelInstance`` / prior instances
+through ``jax.jit`` and ``jax.vmap`` directly.
+"""
+
+from autoconf.jax_wrapper import register_pytree_node, register_pytree_node_class
+
+_ENABLED = False
+_REGISTERED_INSTANCE_CLASSES: set = set()
+
+
+def enable_pytrees() -> bool:
+    """Register PyAutoFit model and prior classes as JAX pytree nodes.
+
+    Returns ``True`` if registration was performed, ``False`` if JAX is not
+    installed (in which case the call is a silent no-op). Safe to call more
+    than once: subsequent calls return ``True`` without re-registering.
+    """
+    global _ENABLED
+    if _ENABLED:
+        return True
+
+    try:
+        import jax  # noqa: F401
+    except ImportError:
+        return False
+
+    import autofit as af
+    from autofit.mapper.prior_model.prior_model import Model
+    from autofit.mapper.prior_model.collection import Collection
+
+    for cls in (
+        af.GaussianPrior,
+        af.UniformPrior,
+        af.LogGaussianPrior,
+        af.LogUniformPrior,
+        af.TruncatedGaussianPrior,
+        Model,
+        Collection,
+        af.ModelInstance,
+    ):
+        try:
+            register_pytree_node_class(cls)
+        except ValueError:
+            # Already registered (e.g. by another caller or a test fixture).
+            # Re-registration is a JAX error but harmless for our purposes.
+            pass
+
+    _ENABLED = True
+    return True
+
+
+def register_model(model) -> bool:
+    """Register every concrete ``model.cls`` in ``model`` as a JAX pytree node.
+
+    PyAutoFit's ``Model.instance_flatten`` / ``instance_unflatten`` produce
+    flatten/unflatten functions for instances of the user-defined class
+    referenced by ``Model.cls`` (e.g. a ``Galaxy`` or ``Gaussian``). For JAX
+    to recurse into those instances rather than treat them as opaque leaves,
+    the user class itself must be registered with ``jax.tree_util``.
+
+    This walks ``model`` (a ``Model`` or ``Collection``) and registers each
+    ``cls`` it finds. Re-registering the same class is a silent no-op. Returns
+    ``True`` if registration ran (or was already complete), ``False`` if JAX
+    is missing.
+    """
+    if not enable_pytrees():
+        return False
+
+    from autofit.mapper.prior_model.prior_model import Model
+    from autofit.mapper.prior_model.collection import Collection
+
+    def _walk(node):
+        if isinstance(node, Model):
+            cls = node.cls
+            if cls not in _REGISTERED_INSTANCE_CLASSES:
+                flatten, unflatten = _build_instance_pytree_funcs(node)
+                try:
+                    register_pytree_node(cls, flatten, unflatten)
+                except ValueError:
+                    # Already registered elsewhere — keep going.
+                    pass
+                _REGISTERED_INSTANCE_CLASSES.add(cls)
+            for _, sub in node.direct_prior_model_tuples:
+                _walk(sub)
+        elif isinstance(node, Collection):
+            for _, sub in node.direct_prior_model_tuples:
+                _walk(sub)
+
+    _walk(model)
+    return True
+
+
+def _build_instance_pytree_funcs(model):
+    """Build flatten/unflatten functions for instances of ``model.cls``.
+
+    Constants from the original model definition (e.g. ``Galaxy(redshift=0.5)``)
+    are placed in the JAX ``aux_data`` so they remain concrete Python values
+    inside a ``jax.jit`` trace. Only prior-derived constructor arguments are
+    placed in ``children`` (and therefore become JAX tracers).
+
+    This is critical for code that uses constants for control flow — e.g.
+    ``sorted(galaxies, key=lambda g: g.redshift)`` in ``Tracer`` — which would
+    otherwise raise ``TracerBoolConversionError`` under JIT.
+    """
+    constructor_args = list(model.constructor_argument_names)
+    constant_arg_names = [
+        name for name in constructor_args if name in dict(model.direct_instance_tuples)
+    ]
+    constant_values = {
+        name: dict(model.direct_instance_tuples)[name] for name in constant_arg_names
+    }
+    dynamic_arg_names = [
+        name for name in constructor_args if name not in constant_arg_names
+    ]
+
+    def flatten(instance):
+        attribute_names = [
+            name
+            for name in model.direct_argument_names
+            if hasattr(instance, name) and name not in constructor_args
+        ]
+        children = (
+            [getattr(instance, name) for name in dynamic_arg_names],
+            [getattr(instance, name) for name in attribute_names],
+        )
+        aux = (
+            tuple(dynamic_arg_names),
+            tuple(constant_arg_names),
+            tuple(constant_values[n] for n in constant_arg_names),
+            tuple(attribute_names),
+        )
+        return children, aux
+
+    def unflatten(aux, children):
+        dyn_names, const_names, const_vals, attr_names = aux
+        dyn_vals, attr_vals = children
+        kwargs = dict(zip(dyn_names, dyn_vals))
+        kwargs.update(zip(const_names, const_vals))
+        ordered = [kwargs[name] for name in constructor_args]
+        instance = model.cls(*ordered)
+        for name, value in zip(attr_names, attr_vals):
+            setattr(instance, name, value)
+        return instance
+
+    return flatten, unflatten

test_autofit/jax/test_enable_pytrees.py

@@ -0,0 +1,101 @@
+"""Tests for ``autofit.jax.enable_pytrees`` / ``register_model``."""
+import pytest
+
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
+
+import autofit as af
+from autofit.jax import enable_pytrees, register_model
+
+
+@pytest.fixture(name="model")
+def make_model():
+    return af.Model(
+        af.ex.Gaussian,
+        centre=af.GaussianPrior(mean=1.0, sigma=1.0),
+        normalization=af.GaussianPrior(mean=2.0, sigma=1.0),
+        sigma=af.GaussianPrior(mean=3.0, sigma=1.0),
+    )
+
+
+def test_enable_pytrees_returns_true_when_jax_available():
+    assert enable_pytrees() is True
+
+
+def test_register_model_round_trip(model):
+    register_model(model)
+    instance = model.instance_from_prior_medians()
+
+    leaves, treedef = jax.tree_util.tree_flatten(instance)
+    rebuilt = jax.tree_util.tree_unflatten(treedef, leaves)
+
+    assert isinstance(rebuilt, af.ex.Gaussian)
+    assert float(rebuilt.centre) == float(instance.centre)
+    assert float(rebuilt.normalization) == float(instance.normalization)
+    assert float(rebuilt.sigma) == float(instance.sigma)
+
+
+def test_register_model_works_under_jit(model):
+    register_model(model)
+    instance = model.instance_from_prior_medians()
+
+    @jax.jit
+    def total(inst):
+        return inst.centre + inst.normalization + inst.sigma
+
+    assert float(total(instance)) == pytest.approx(
+        instance.centre + instance.normalization + instance.sigma
+    )
+
+
+def test_register_model_keeps_constants_static():
+    """Constants from the model definition must NOT be traced under JIT.
+
+    Galaxy.redshift is a common case: it sits in the constructor signature
+    but gets a fixed value via ``af.Constant``. If it became a JAX tracer,
+    code paths like ``sorted(galaxies, key=lambda g: g.redshift)`` would
+    blow up with ``TracerBoolConversionError``. The ``aux_data`` partition
+    in ``register_model`` keeps it as a Python float.
+    """
+    class Holder:
+        def __init__(self, redshift, scale):
+            self.redshift = redshift
+            self.scale = scale
+
+    model = af.Model(
+        Holder,
+        redshift=0.5,
+        scale=af.GaussianPrior(mean=1.0, sigma=1.0),
+    )
+    register_model(model)
+    instance = model.instance_from_prior_medians()
+    assert isinstance(instance.redshift, float)
+
+    @jax.jit
+    def use_redshift_for_control_flow(inst):
+        # `if` on a traced value would raise TracerBoolConversionError;
+        # this only works if redshift is kept static via aux_data.
+        if inst.redshift > 0:
+            return inst.scale * 2
+        return inst.scale
+
+    result = use_redshift_for_control_flow(instance)
+    assert float(result) == pytest.approx(2.0 * instance.scale)
+
+
+def test_enable_pytrees_idempotent():
+    assert enable_pytrees() is True
+    assert enable_pytrees() is True
+
+
+def test_collection_round_trip(model):
+    register_model(model)
+    collection = af.Collection(g1=model, g2=model)
+    register_model(collection)
+
+    instance = collection.instance_from_prior_medians()
+    leaves, treedef = jax.tree_util.tree_flatten(instance)
+    rebuilt = jax.tree_util.tree_unflatten(treedef, leaves)
+
+    assert float(rebuilt.g1.centre) == float(instance.g1.centre)
+    assert float(rebuilt.g2.sigma) == float(instance.g2.sigma)