jax: gated pytree registration for AbstractNDArray + register_instance_pytree helper

Reintroduces the removed auto-registration of AbstractNDArray subclasses as
JAX pytrees, gated so it only runs when an instance is constructed on the
JAX path (xp is not np). Each subclass pays the registration cost at most
once via a module-level sentinel.

Also adds register_instance_pytree(cls, no_flatten=...), a generic helper
for non-AbstractNDArray classes (FitImaging, Tracer, DatasetModel, etc.)
that flattens __dict__ and carries no_flatten names through aux_data for
per-analysis constants like dataset/settings/cosmology.

Issue: PyAutoLabs/PyAutoLens#444

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -62,6 +62,75 @@ def wrapper(self, other):
     return wrapper
 
 
+_pytree_registered_classes: set = set()
+
+
+def _register_as_pytree(cls):
+    """Register ``cls`` with ``jax.tree_util`` via the lazy autoconf wrapper.
+
+    Gated: only called when a subclass instance is constructed on the JAX path
+    (``xp is not np``). The registration is class-scoped via
+    ``_pytree_registered_classes`` so each subclass pays the cost at most once
+    regardless of how many instances are created. No-ops if JAX is not installed.
+    """
+    if cls in _pytree_registered_classes:
+        return
+    from autoconf.jax_wrapper import register_pytree_node
+
+    register_pytree_node(cls, cls.instance_flatten, cls.instance_unflatten)
+    _pytree_registered_classes.add(cls)
+
+
+def register_instance_pytree(cls, no_flatten=()):
+    """Register any class with ``jax.tree_util`` via ``__dict__`` flattening.
+
+    Generic counterpart to :func:`_register_as_pytree` for classes that are
+    *not* ``AbstractNDArray`` subclasses but still need to round-trip through
+    ``jax.jit`` (e.g. ``FitImaging``, ``Tracer``, ``Imaging``). Attributes are
+    partitioned using ``no_flatten``:
+
+    * Names **not** in ``no_flatten`` ride as pytree children — JAX traces them
+      and can substitute new values on unflatten (dynamic per fit).
+    * Names **in** ``no_flatten`` ride as ``aux_data`` — JAX treats them as
+      opaque Python objects, closing over the original reference across the
+      JIT boundary. Appropriate for per-analysis constants (dataset, settings,
+      cosmology, adapt images).
+
+    Reconstructs via ``cls.__new__`` + ``setattr`` (side-effect-free — no
+    ``__init__`` re-entry). Idempotent.
+    """
+    if cls in _pytree_registered_classes:
+        return
+    from autoconf.jax_wrapper import register_pytree_node
+
+    no_flatten_set = frozenset(no_flatten)
+
+    def flatten(instance):
+        dyn: list = []
+        static: list = []
+        for key, value in sorted(instance.__dict__.items()):
+            if key in no_flatten_set:
+                static.append((key, value))
+            else:
+                dyn.append((key, value))
+        dyn_keys = tuple(k for k, _ in dyn)
+        dyn_values = tuple(v for _, v in dyn)
+        static_items = tuple(static)
+        return dyn_values, (dyn_keys, static_items)
+
+    def unflatten(aux, children):
+        dyn_keys, static_items = aux
+        new = cls.__new__(cls)
+        for key, value in zip(dyn_keys, children):
+            setattr(new, key, value)
+        for key, value in static_items:
+            setattr(new, key, value)
+        return new
+
+    register_pytree_node(cls, flatten, unflatten)
+    _pytree_registered_classes.add(cls)
+
+
 class AbstractNDArray(ABC):
 
     __no_flatten__ = ()
@@ -76,6 +145,9 @@ def __init__(self, array, xp=np):
 
         self.use_jax = xp is not np
 
+        if self.use_jax:
+            _register_as_pytree(type(self))
+
     @property
     def is_transformed(self) -> bool:
         return self._is_transformed

test_autoarray/test_jax_pytree.py

@@ -0,0 +1,64 @@
+"""Tests for gated JAX pytree registration of ``AbstractNDArray`` subclasses.
+
+Follows the three-step pattern from ``autolens_workspace_test/scripts/hessian_jax.py``:
+1. NumPy path — confirm autoarray type with ``np.ndarray`` backing, no pytree registration.
+2. JAX path outside JIT — same autoarray type with ``jax.Array`` backing; pytree registered.
+3. JAX path through ``jax.jit`` — round-trip the instance and assert the output carries
+   a ``jax.Array`` leaf.
+"""
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
+
+from autoarray.abstract_ndarray import AbstractNDArray, _pytree_registered_classes
+
+
+class _LeafArray(AbstractNDArray):
+    """Minimal concrete ``AbstractNDArray`` with no nested autoarray children.
+
+    Isolates the pytree-registration machinery from the larger autoarray
+    hierarchy: a real ``Array2D`` also carries a ``Mask2D`` and other nested
+    ``AbstractNDArray`` children whose own registration is covered by
+    follow-up steps in the ``fit-imaging-pytree`` task.
+    """
+
+    @property
+    def native(self):
+        return self
+
+
+def test_numpy_path_does_not_register_pytree():
+    _pytree_registered_classes.discard(_LeafArray)
+
+    arr = _LeafArray(np.array([1.0, 2.0, 3.0]))
+
+    assert isinstance(arr._array, np.ndarray)
+    assert _LeafArray not in _pytree_registered_classes
+
+
+def test_jax_path_registers_pytree_once():
+    _pytree_registered_classes.discard(_LeafArray)
+
+    arr_jax = _LeafArray(jnp.array([1.0, 2.0, 3.0]), xp=jnp)
+
+    assert isinstance(arr_jax._array, jnp.ndarray)
+    assert _LeafArray in _pytree_registered_classes
+
+    # Second construction on the JAX path is a no-op; class stays registered.
+    _LeafArray(jnp.array([4.0, 5.0]), xp=jnp)
+    assert _LeafArray in _pytree_registered_classes
+
+
+def test_jax_jit_round_trip_returns_wrapper_with_jax_array():
+    arr_jax = _LeafArray(jnp.array([1.0, 2.0, 3.0]), xp=jnp)
+    assert _LeafArray in _pytree_registered_classes
+
+    result = jax.jit(lambda a: a)(arr_jax)
+
+    assert isinstance(result, _LeafArray)
+    assert isinstance(result._array, jnp.ndarray)
+    npt.assert_allclose(np.asarray(result._array), np.asarray(arr_jax._array))