Merge pull request #421 from PyAutoLabs/feature/simulator-jax-xp-threading

fix: make the imaging simulator's @jax.jit path work (5 xp / pytree sites)

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -164,7 +164,20 @@ def is_transformed(self, value: bool):
 
     @property
     def _xp(self):
-        if self.use_jax:
+        """The array module matching the array this instance actually holds.
+
+        ``use_jax`` records how the instance was *constructed*, and not every
+        construction site threads ``xp`` — an intermediate operation can hand
+        back an instance whose flag says NumPy while its backing array is a JAX
+        tracer. Consumers of ``_xp`` (notably ``Array2D.native``, which is a
+        property and so cannot be passed ``xp`` by its caller) would then route a
+        traced array through the NumPy path and raise
+        ``TracerArrayConversionError`` inside a ``jax.jit`` trace.
+
+        So the backing array is authoritative: if it is a JAX type, this is
+        ``jnp`` regardless of the flag. On the NumPy path nothing changes.
+        """
+        if self.use_jax or type(self._array).__module__.startswith(("jax", "jaxlib")):
             import jax.numpy as jnp
 
             return jnp

autoarray/dataset/imaging/simulator.py

@@ -12,6 +12,57 @@
 
 logger = logging.getLogger(__name__)
 
+_IMAGING_PYTREES_REGISTERED = False
+
+
+def _register_imaging_pytrees() -> None:
+    """Register ``Imaging`` so ``jax.jit(via_image_from)`` can flatten its return.
+
+    Counterpart of ``AnalysisImaging._register_fit_imaging_pytrees``, which does
+    the same job for ``FitImaging``. Without it a jitted simulator call raises
+    ``TypeError: ... returned a value of type Imaging, which is not a valid JAX
+    type``.
+
+    ``data`` and ``noise_map`` are the only per-simulation dynamic values, so
+    everything else rides as aux: ``psf`` and ``grids`` are constants for a given
+    simulation, the over-sample sizes are static integer geometry, and the
+    remaining slots are ``None``.
+
+    Unlike pytree registration for a jitted function's *arguments* — which must
+    happen before the first call, because JAX flattens arguments at trace time —
+    registering here is in time, because the return value is flattened only after
+    the body has run. Idempotent via the module-level flag.
+    """
+    global _IMAGING_PYTREES_REGISTERED
+    if _IMAGING_PYTREES_REGISTERED:
+        return
+
+    from autoarray.abstract_ndarray import register_instance_pytree
+
+    # ``Array2D.instance_flatten`` emits every ``__dict__`` entry not opted out,
+    # so a flattened ``data`` exposes its ``mask`` as a child. ``Mask2D`` must
+    # therefore be a pytree itself, or it surfaces as a bare leaf and JAX
+    # rejects the return value. Registering it here rather than adding ``mask``
+    # to ``Array2D.__no_flatten__`` keeps ``Array2D``'s flatten semantics
+    # unchanged for every other jitted path in the stack.
+    register_instance_pytree(Mask2D)
+
+    register_instance_pytree(
+        Imaging,
+        no_flatten=(
+            "psf",
+            "grids",
+            "over_sample_size_lp",
+            "over_sample_size_pixelization",
+            "convolve_over_sample_size_lp",
+            "convolve_over_sample_size_pixelization",
+            "noise_covariance_matrix",
+            "sparse_operator",
+        ),
+    )
+
+    _IMAGING_PYTREES_REGISTERED = True
+
 
 class SimulatorImaging:
     def __init__(
@@ -162,6 +213,9 @@ def via_image_from(
         if xp is None:
             xp = self._xp
 
+        if xp is not np:
+            _register_imaging_pytrees()
+
         exposure_time_map = Array2D.full(
             fill_value=self.exposure_time,
             shape_native=image.shape_native,
@@ -201,7 +255,9 @@ def via_image_from(
 
         if self.include_poisson_noise_in_noise_map:
             noise_map = preprocess.noise_map_via_data_eps_and_exposure_time_map_from(
-                data_eps=image_with_poisson_noise, exposure_time_map=exposure_time_map
+                data_eps=image_with_poisson_noise,
+                exposure_time_map=exposure_time_map,
+                xp=xp,
             )
 
         else:

autoarray/dataset/preprocess.py

@@ -132,7 +132,9 @@ def array_with_random_uniform_values_added(array, upper_limit=0.001):
     return array + upper_limit * np.random.uniform(size=array.shape_slim)
 
 
-def noise_map_via_data_eps_and_exposure_time_map_from(data_eps, exposure_time_map):
+def noise_map_via_data_eps_and_exposure_time_map_from(
+    data_eps, exposure_time_map, xp=np
+):
     """
     Estimate the noise-map value in every data-point, by converting the data to units of counts and taking the
     square root of these values.
@@ -148,9 +150,13 @@ def noise_map_via_data_eps_and_exposure_time_map_from(data_eps, exposure_time_ma
         The data in electrons second used to estimate the Poisson noise in every data point.
     exposure_time_map
         The exposure time at every data-point of the data.
+    xp
+        The array module (``numpy`` or ``jax.numpy``). Must be ``jnp`` when
+        ``data_eps`` carries a traced array, otherwise ``np.abs`` raises
+        ``TracerArrayConversionError`` inside a ``jax.jit`` trace.
     """
     return data_eps.with_new_array(
-        np.abs(data_eps.array * exposure_time_map.array) ** 0.5
+        xp.abs(data_eps.array * exposure_time_map.array) ** 0.5
         / exposure_time_map.array
     )
 

autoarray/structures/arrays/array_2d_util.py

@@ -509,8 +509,15 @@ def array_2d_native_from(
 
     shape = (mask_2d.shape[0], mask_2d.shape[1])
 
+    # Deliberately NumPy, not ``xp``: this index map is derived from the mask,
+    # which is concrete geometry and never traced. Its computation is a
+    # ``where(~mask)``, whose output shape depends on the *values* of the mask, so
+    # under ``jnp`` inside a ``jax.jit`` trace it raises
+    # ``ConcretizationTypeError``. Computing it in NumPy yields a concrete index
+    # array, which is a valid static operand for the scatter below — that is
+    # where ``xp`` genuinely belongs.
     native_index_for_slim_index_2d = mask_2d_util.native_index_for_slim_index_2d_from(
-        mask_2d=mask_2d, xp=xp
+        mask_2d=mask_2d
     ).astype("int")
 
     return array_2d_via_indexes_from(

autoarray/structures/arrays/uniform_2d.py

@@ -291,9 +291,18 @@ def native(self) -> "Array2D":
 
         If it is already stored in its `native` representation it is return as it is. If not, it is mapped from
         `slim` to `native` and returned as a new `Array2D`.
+
+        ``xp`` is forwarded from ``self._xp`` rather than left at its ``numpy``
+        default. A property cannot take an argument, so without this a JAX-backed
+        array would be re-mapped slim -> native through the NumPy path and raise
+        ``TracerArrayConversionError`` inside a ``jax.jit`` trace.
         """
         return Array2D(
-            values=self, mask=self.mask, header=self.header, store_native=True
+            values=self,
+            mask=self.mask,
+            header=self.header,
+            store_native=True,
+            xp=self._xp,
         )
 
     @property