fix: make array_2d_via_indexes_from jit-traceable

Replace tuple(native_index_for_slim_index_2d.T) with 2D advanced
indexing on the JAX path — the tuple() call iterates the outermost
axis of a traced array, triggering TracerArrayConversionError under
@jax.jit. Remove the temporary if/else workaround in the imaging
simulator and the "jit blocked" docstring caveats.

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

Author: 2 people

autoarray/dataset/imaging/simulator.py

@@ -95,14 +95,6 @@ def __init__(
             The returned ``Imaging`` carries ``jax.Array`` data — useful for
             JAX-eager batched simulation (parameter sweeps, mock-data studies).
             Defaults to ``False``.
-
-            Note: ``@jax.jit`` wrapping of ``via_tracer_from`` /
-            ``via_galaxies_from`` is currently blocked by an unrelated
-            pre-existing limitation — ``Array2D.native`` is not jit-traceable
-            because it goes through indexed assignment in
-            ``array_2d_via_indexes_from``. A separate PyAutoArray task is
-            needed to refactor the slim/native reshape to be jit-friendly.
-            Eager JAX usage works today.
         """
 
         if psf is not None:
@@ -232,16 +224,7 @@ def via_image_from(
             origin=image.origin,
         )
 
-        # Re-wrap the image against the all-false mask. Use ``.array`` rather
-        # than ``.native`` on the JAX path: ``.native`` routes through
-        # ``array_2d_via_indexes_from`` which is not jit-safe (it builds a
-        # native-shape array by Python-iterating an index tuple). ``.array``
-        # returns the raw backing array which the ``Array2D`` constructor
-        # accepts in either slim or native shape.
-        if xp is np:
-            image = Array2D(values=image.native, mask=mask)
-        else:
-            image = Array2D(values=image.array, mask=mask)
+        image = Array2D(values=image.native, mask=mask)
 
         dataset = Imaging(
             data=image,

autoarray/dataset/interferometer/simulator.py

@@ -64,10 +64,7 @@ def __init__(
             If ``True``, ``via_image_from`` defaults ``xp`` to ``jax.numpy`` and
             the simulator's internal complex-Gaussian noise generation routes
             through ``jax.random``. The returned ``Interferometer`` carries
-            ``jax.Array`` visibilities. Mirror of ``SimulatorImaging.use_jax``;
-            same caveat applies — ``@jax.jit`` wrapping is currently blocked
-            by autoarray's pre-existing ``.native`` reshape limitation in the
-            transformer / dataset construction path. Eager JAX usage works.
+            ``jax.Array`` visibilities. Mirror of ``SimulatorImaging.use_jax``.
         """
 
         self.uv_wavelengths = uv_wavelengths

autoarray/structures/arrays/array_2d_util.py

@@ -557,7 +557,10 @@ def array_2d_via_indexes_from(
     array = xp.zeros(shape, dtype=array_2d_slim.dtype)
 
     if xp.__name__.startswith("jax"):
-        array = array.at[tuple(native_index_for_slim_index_2d.T)].set(array_2d_slim)
+        array = array.at[
+            native_index_for_slim_index_2d[:, 0],
+            native_index_for_slim_index_2d[:, 1],
+        ].set(array_2d_slim)
     else:
         array[tuple(native_index_for_slim_index_2d.T)] = array_2d_slim