Affected Repositories

• PyAutoArray (primary)
• autolens_workspace_test
• autolens_workspace
• autogalaxy_workspace

Work Classification

Both (library first, workspace follow-up)

Branch Survey

Suggested branch: feature/array2d-native-jit-safety
Worktree root: ~/Code/PyAutoLabs-wt/array2d-native-jit-safety/ (created later by /start_library)

Implementation Steps

1. Fix array_2d_via_indexes_from in autoarray/structures/arrays/array_2d_util.py:557-564:
   Replace array.at[tuple(native_index_for_slim_index_2d.T)].set(...) with array.at[native_index_for_slim_index_2d[:, 0], native_index_for_slim_index_2d[:, 1]].set(...) — 2D advanced indexing instead of tuple-unpacked rows.

2. Remove imaging simulator workaround in autoarray/dataset/imaging/simulator.py:235-244:
   Remove the if xp is np: image.native else: image.array guard and its comment, restore the simple Array2D(values=image.native, mask=mask).

3. Update simulator docstrings:

   • autoarray/dataset/imaging/simulator.py:99-105 — remove "Note: @jax.jit wrapping ... currently blocked" paragraph
   • autoarray/dataset/interferometer/simulator.py:68-70 — remove "@jax.jit wrapping is currently blocked" sentence

4. Restore workspace_test parity scripts:

   • autolens_workspace_test/scripts/imaging/simulator_use_jax_parity.py:96-114 — replace commented-out JIT block with active @jax.jit roundtrip test
   • autolens_workspace_test/scripts/interferometer/simulator_use_jax_parity.py — update docstring + add JIT roundtrip test

5. Workspace doc cleanup — remove "blocked by Array2D.native" notes from JAX variant blocks in:

   • autolens_workspace/scripts/imaging/simulator.py
   • autolens_workspace/scripts/interferometer/simulator.py
   • autolens_workspace/scripts/group/simulator.py
   • autolens_workspace/scripts/point_source/simulator.py
   • autolens_workspace/scripts/point_source/fit.py
   • autogalaxy_workspace/scripts/imaging/simulator.py
   • autogalaxy_workspace/scripts/interferometer/simulator.py

Key Files

• autoarray/structures/arrays/array_2d_util.py — core fix (the only tuple(...T) pattern on traced arrays)
• autoarray/dataset/imaging/simulator.py — workaround removal + docstring
• autoarray/dataset/interferometer/simulator.py — docstring
• autolens_workspace_test/scripts/imaging/simulator_use_jax_parity.py — restore JIT test
• autolens_workspace_test/scripts/interferometer/simulator_use_jax_parity.py — restore JIT test

Refactor Array2D.native / array_2d_via_indexes_from for JAX-jit safety

SimulatorImaging.use_jax=True and SimulatorInterferometer.use_jax=True
work end-to-end on JAX-eager today (shipped in Phase 2 of
z_features/jax_user_intro.md — PRs PyAutoArray#335, #336 +
PyAutoLens#539, #540 + PyAutoGalaxy#442, #443) but cannot currently
be wrapped in @jax.jit. This task unblocks that.

The blocker

When the user wraps a simulator in their own @jax.jit:

simulator = al.SimulatorImaging(..., use_jax=True)

@jax.jit
def simulate(tracer):
    return simulator.via_tracer_from(tracer=tracer, grid=grid)

The call fails with TracerArrayConversionError: __array__() was called on traced array. The two failing sites are:

1. array_2d_via_indexes_from (array_2d_util.py:535-562) — tuple(jax_array.T) iterates the outermost axis, forcing __array__().
2. Array2D.native (uniform_2d.py:295) — routes through (1).
3. Imaging.trimmed_after_convolution_from — same .native access.

Scope

• Refactor array_2d_via_indexes_from to use JAX-traceable 2D advanced indexing
• Remove the temporary if xp is np simulator guard
• Update simulator docstrings to drop the "jit currently blocked" caveat
• Restore the disabled @jax.jit tests in workspace_test parity scripts
• Remove "blocked by Array2D.native" notes from workspace JAX variant blocks