fix: simulator @jax.jit path fails on un-threaded xp sites (preprocess + array_2d_util)

## Overview

The `@jax.jit` simulator recipe documented in both workspaces does not run. Two
independent causes; **this issue is the second one, which is a real library
defect.**

The first cause was a docs error and is already fixed (registration is the
caller's responsibility and cannot be automated — JAX flattens jitted arguments
at trace time; see PyAutoLabs/autolens_workspace#379, merged). But *even with*
correct pytree registration, the jitted call still fails inside autoarray on
array sites that do not thread `xp`.

Reproduced 2026-07-30 in **both** libraries — identical failure and identical
frame, so this is autoarray, not autolens/autogalaxy:

```python
register_tracer_classes(tracer)          # or register_galaxies_classes(galaxies)

simulator = al.SimulatorImaging(exposure_time=300.0, psf=psf,
                                background_sky_level=0.1, use_jax=True)

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

simulate(tracer)
# jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method
# __array__() was called on traced array with shape float64[3600]
```

### Site 1 — `preprocess.py:153`

Frame chain:

```
autolens/imaging/simulator.py:91            via_tracer_from -> via_image_from
autoarray/dataset/imaging/simulator.py:203  via_image_from
autoarray/dataset/preprocess.py:153         noise_map_via_data_eps_and_exposure_time_map_from
```

`noise_map_via_data_eps_and_exposure_time_map_from` (preprocess.py:134-155)
**takes no `xp` parameter at all** and hardcodes NumPy:

```python
return data_eps.with_new_array(
    np.abs(data_eps.array * exposure_time_map.array) ** 0.5
    / exposure_time_map.array
)
```

Its call site cannot pass `xp` because the parameter does not exist.
`include_poisson_noise_in_noise_map` defaults to **`True`**, so the default path
always hits it. (Toggling `add_poisson_noise_to_data` does not avoid it — that is
a different flag.)

### Site 2 — `array_2d_util.py`

Bypassing site 1 with `include_poisson_noise_in_noise_map=False` moves the
failure into the `Array2D.full(...)` else-branch:

```
autoarray/structures/arrays/array_2d_util.py:147  convert_array_2d
autoarray/structures/arrays/array_2d_util.py:516  array_2d_native_from
autoarray/structures/arrays/array_2d_util.py:565  array_2d_via_indexes_from
```

**Depth beyond these two is unmeasured** — expect to iterate: fix a site, re-run,
find the next.

### Related observation (confirm before treating as a defect)

Eager `via_tracer_from` with `use_jax=True` **succeeds** but returns
`dataset.data.array` of type `numpy.ndarray`, not `jax.Array`, with
`add_poisson_noise_to_data` either `True` or `False`. The `use_jax` plumbing is
real and not a no-op kwarg (`autoarray/dataset/imaging/simulator.py:115-122` sets
`self.use_jax`, `_xp` returns `jnp`, and `via_tracer_from` threads `xp`), so the
NumPy backing is plausibly a symptom of these same un-threaded sites falling back
to `np`. Worth confirming whether eager output is *supposed* to be JAX-backed
before treating it as separate.

## Plan

- Thread `xp` through the autoarray sites the jitted simulator path reaches,
  starting with the two identified, iterating until the recipe runs and returns a
  `jax.Array`-backed dataset.
- Add `xp=np` **default** parameters — do not change existing defaults or
  behaviour on the NumPy path.
- Re-check the eager return-type question once the sites are threaded.

<details>
<summary>Detailed implementation plan</summary>

### Work Classification

Library (PyAutoArray)

### Affected Repositories

- PyAutoArray (primary)

### Branch Survey

| Repository | Current Branch | Dirty? |
|-----------|---------------|--------|
| ./PyAutoArray | main | clean |

**Suggested branch:** `feature/simulator-jax-xp-threading`

### Steps

1. `autoarray/dataset/preprocess.py` — add `xp=np` to
   `noise_map_via_data_eps_and_exposure_time_map_from`, replace `np.abs` with
   `xp.abs`; thread `xp=xp` from `autoarray/dataset/imaging/simulator.py:202-204`.
2. Re-run the reproducer; fix whatever `array_2d_util.py` site blocks next
   (`convert_array_2d` / `array_2d_native_from` / `array_2d_via_indexes_from`),
   reached via the `Array2D.full(...)` else-branch.
3. Repeat until `jax.jit(simulate)(tracer)` returns an `Imaging` whose
   `.data.array` is a `jax.Array`.
4. Check the interferometer path (`via_image_from` → `Interferometer`) for the
   same class of site.

### Testing

Unit tests are NumPy-only across the stack, so this cannot be covered by
`test_autoarray/`. Validate by:

- `python -m pytest test_autoarray/` stays green (additive `xp=np` defaults).
- Reproducer: registration + `jax.jit(simulate)(tracer)` returns a
  `jax.Array`-backed `Imaging`, in **both** autolens and autogalaxy.
- Jitted result matches the eager NumPy simulation for a fixed `noise_seed`.
- Both `include_poisson_noise_in_noise_map` branches work under jit.
- Interferometer equivalent (`TransformerDFT` and the nufftax-backed
  `TransformerNUFFT`; `TransformerNUFFTPyNUFFT` is not JAX-traceable).
- Downstream: `pytest test_autogalaxy/` and `test_autolens/`.

### Docs follow-up (separate, after this merges)

`autolens_workspace/scripts/guides/using_jax.py`,
`autogalaxy_workspace/scripts/guides/using_jax.py` and the four
`simulator.py` `__JAX Variant__` sections currently say the jitted wrap "does not
currently work" and link here. Once fixed, restore the recipe **and uncomment the
call** so CI executes it — the governing lesson from #368/#379 is that a
documented JAX recipe no script executes will be wrong.

</details>

## Original Prompt

<details>
<summary>Click to expand starting prompt</summary>

# The documented `@jax.jit` simulator recipe does not run — un-threaded `xp` sites in autoarray

Type: bug
Target: autoarray
Repos:
- PyAutoArray
Difficulty: medium
Autonomy: supervised
Priority: normal

Both workspaces' `scripts/guides/using_jax.py` and the `__JAX Variant__` sections
of five `simulator.py` scripts document this recipe:

```python
simulator = al.SimulatorImaging(
    exposure_time=300.0, psf=psf, background_sky_level=0.1, use_jax=True
)

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

with the claim *"The simulator handles pytree registration internally — you write
nothing JAX-specific beyond the decorator."*

**Every part of that is false, and the recipe does not run.** Measured
2026-07-29 on the installed stack (`autolens` at `main`), 50×50 grid, Gaussian
PSF via `al.Convolver.from_gaussian`.

## Finding 1 — nothing registers pytrees, anywhere

`grep -rn register` over `autolens/imaging/simulator.py`,
`autolens/interferometer/simulator.py`, `autogalaxy/imaging/simulator.py`,
`autogalaxy/interferometer/simulator.py` and both autoarray simulator bases
(`autoarray/dataset/{imaging,interferometer}/simulator.py`): **zero hits**. No
simulator calls `register_tracer_classes` or `register_galaxies_pytree`.

Without registration the jitted call fails immediately:

```
TypeError: Error interpreting argument to <function simulate ...> as an abstract
array. The problematic value is of type <class 'autolens.lens.tracer.Tracer'>
and was passed to the function at path tracer.
```

`PointSolver.solve_triangles` already documents **why** a simulator cannot fix
this itself (`autolens/point/solver/point_solver.py:102-109`):

> NOTE: pytree registration is the user's responsibility ... Auto-registering
> inside `solve()` doesn't help because JAX flattens function arguments at trace
> time — before entering this method — so registration must run before the first
> jitted call.

That constraint applies identically to `via_tracer_from`. So the claim is not
merely unimplemented, it is **unimplementable as worded** — the user must call
`register_tracer_classes(tracer)` (autolens) / `register_galaxies_classes(galaxies)`
(autogalaxy, added in PyAutoGalaxy#537) before the first jitted call.

Note `autolens/jax/registration.py`'s own module docstring is stale on the same
point: it claims registration "is called automatically by
`PointSolver(use_jax=True).solve(tracer, ...)` on the first invocation and by
`Simulator(use_jax=True).via_tracer_from(tracer, ...)` ... once Phase 2 ships the
Simulator changes". Neither is true; nothing in autolens calls
`register_tracer_classes` at all.

## Finding 2 — even WITH correct registration, the jitted call still fails

This is the library bug. After `register_tracer_classes(tracer)`:

```
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method
__array__() was called on traced array with shape float64[3600]
```

Failing frame chain:

```
autolens/imaging/simulator.py:91            via_tracer_from  ->  via_image_from
autoarray/dataset/imaging/simulator.py:203  via_image_from
autoarray/dataset/preprocess.py:153         noise_map_via_data_eps_and_exposure_time_map_from
```

`noise_map_via_data_eps_and_exposure_time_map_from` (preprocess.py:134-155)
**takes no `xp` parameter at all** and hardcodes NumPy:

```python
return data_eps.with_new_array(
    np.abs(data_eps.array * exposure_time_map.array) ** 0.5
    / exposure_time_map.array
)
```

Its call site (`autoarray/dataset/imaging/simulator.py:202-204`) cannot pass `xp`
because the parameter does not exist. `include_poisson_noise_in_noise_map`
defaults to **`True`**, so the default path always hits it. Toggling
`add_poisson_noise_to_data` does not avoid it — that is a different flag.

**At least two independent sites.** Bypassing the first with
`include_poisson_noise_in_noise_map=False` moves the failure to the
`Array2D.full(...)` else-branch:

```
autoarray/structures/arrays/array_2d_util.py:147  convert_array_2d
autoarray/structures/arrays/array_2d_util.py:516  array_2d_native_from
autoarray/structures/arrays/array_2d_util.py:565  array_2d_via_indexes_from
```

So this is a multi-site `xp`-threading job, not a one-line fix, and the depth
beyond these two is unmeasured — expect to iterate: fix a site, re-run, find the
next.

## Finding 3 — `use_jax=True` does not put the EAGER result on JAX either

The simulator scripts also claim *"eager `simulator_jax.via_tracer_from(tracer,
grid)` (no `@jax.jit`) already runs on JAX and is sufficient for one-off
simulations."*

Measured: eager `via_tracer_from` with `use_jax=True` **succeeds** but returns
`dataset.data.array` of type **`numpy.ndarray`**, not `jax.Array` — with
`add_poisson_noise_to_data` either `True` or `False`. The `use_jax` plumbing is
present and not a no-op kwarg (`autoarray/dataset/imaging/simulator.py:115-122`
sets `self.use_jax` and `_xp` returns `jnp`; `via_tracer_from` threads `xp`
through), so the NumPy backing is presumably a symptom of the same un-threaded
sites silently falling back to `np`. Confirm whether eager output is *supposed*
to be JAX-backed before treating this as a separate defect — it may be the
correct consequence of `with_new_array` on a NumPy-produced noise map.

## Why nothing caught this

The `__JAX Variant__` sections leave the actual call commented out —
`autolens_workspace/scripts/imaging/simulator.py:381` is
`# dataset_jax = simulate(tracer)`. So the recipe is prose plus a `@jax.jit`
decorator that is never invoked. Same failure mode as the six
`likelihood_function.py` `__JAX__` blocks removed in autolens_workspace#368.

## Scope

**In scope (this prompt):** thread `xp` through the autoarray sites the simulator
`@jax.jit` path reaches, starting with the two identified, iterating until the
recipe runs and returns a `jax.Array`-backed `Imaging`. Add `xp=np` default
parameters rather than changing defaults.

**Out of scope / separate:**

- **The docs are wrong TODAY and should be corrected regardless of whether this
  bug is fixed** — see the docs follow-up below. Do not leave the false claims
  standing while this bug waits.
- Any change to `register_tracer_classes` / `register_galaxies_classes` auto-call
  behaviour. Per `PointSolver`'s note, auto-registration inside the call cannot
  work; registration stays the user's responsibility and the docs must say so.

## Validation

Unit tests are NumPy-only across the stack, so this cannot be covered by
`test_autoarray/`. Validate with a parity script in `autolens_workspace_test` /
`autogalaxy_workspace_test`:

1. `register_tracer_classes(tracer)` then `jax.jit(simulate)(tracer)` returns an
   `Imaging` whose `.data.array` is a `jax.Array`.
2. The jitted result matches the eager NumPy simulation for a fixed
   `noise_seed`.
3. Both `include_poisson_noise_in_noise_map` branches work under jit.
4. Interferometer equivalent (`via_tracer_from` → `FitInterferometer` path,
   `TransformerDFT` only — `TransformerNUFFT` is not JAX-traceable).
5. `test_autoarray/` stays green (additive `xp=np` defaults).

## Companion docs follow-up (file separately, do first)

Correct the false claims in:

- `autolens_workspace/scripts/guides/using_jax.py` — `__Writing @jax.jit Yourself__` item 1
- `autogalaxy_workspace/scripts/guides/using_jax.py` — `__Writing @jax.jit Yourself__`
- `autolens_workspace/scripts/{imaging,interferometer,point_source}/simulator.py` — `__JAX Variant (Advanced)__`
- `autogalaxy_workspace/scripts/{imaging,interferometer}/simulator.py` — `__JAX Variant__`
- also check `autolens_workspace/scripts/{group,multi_galaxy}/simulator.py` (`__JAX Variant__`) and `cluster/simulator.py` (`__JAX JIT__`, `__JAX JIT — Point Solver__`)
- `autolens/jax/registration.py` module docstring (stale auto-call claims) — library, tiny

Either state the registration requirement and mark the jit recipe as not yet
supported, or uncomment the call so CI proves whatever the docs claim. The
governing lesson from autolens_workspace#368: **a documented JAX recipe that no
script executes will be wrong.**

</details>


Navigation Menu

Uh oh!

fix: simulator @jax.jit path fails on un-threaded xp sites (preprocess + array_2d_util) #420

Overview

Site 1 — `preprocess.py:153`

Site 2 — `array_2d_util.py`

Related observation (confirm before treating as a defect)

Plan

Work Classification

Affected Repositories

Branch Survey

Steps

Testing

Docs follow-up (separate, after this merges)

Original Prompt

The documented `@jax.jit` simulator recipe does not run — un-threaded `xp` sites in autoarray

Finding 1 — nothing registers pytrees, anywhere

Finding 2 — even WITH correct registration, the jitted call still fails

Finding 3 — `use_jax=True` does not put the EAGER result on JAX either

Why nothing caught this

Scope

Validation

Companion docs follow-up (file separately, do first)

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Uh oh!

fix: simulator @jax.jit path fails on un-threaded xp sites (preprocess + array_2d_util) #420

Description

Overview

Site 1 — preprocess.py:153

Site 2 — array_2d_util.py

Related observation (confirm before treating as a defect)

Plan

Work Classification

Affected Repositories

Branch Survey

Steps

Testing

Docs follow-up (separate, after this merges)

Original Prompt

The documented @jax.jit simulator recipe does not run — un-threaded xp sites in autoarray

Finding 1 — nothing registers pytrees, anywhere

Finding 2 — even WITH correct registration, the jitted call still fails

Finding 3 — use_jax=True does not put the EAGER result on JAX either

Why nothing caught this

Scope

Validation

Companion docs follow-up (file separately, do first)

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Site 1 — `preprocess.py:153`

Site 2 — `array_2d_util.py`

The documented `@jax.jit` simulator recipe does not run — un-threaded `xp` sites in autoarray

Finding 3 — `use_jax=True` does not put the EAGER result on JAX either