fix: make the interferometer simulator's @jax.jit path work

Follow-up to #420/#421, which fixed the imaging simulator and deliberately
excluded the interferometer. Two causes, and one of them was not what #422
predicted.

1. Interferometer was not a registered pytree, so it could not cross the jit
   RETURN boundary. Added _register_interferometer_pytrees, mirroring
   _register_imaging_pytrees: `data` and `noise_map` dynamic; uv_wavelengths,
   real_space_mask, transformer, grids, the over-sample sizes and the Nones as
   aux. It must also register Visibilities and VisibilitiesNoiseMap — unlike
   Array2D on the imaging path, nothing else registers those, so they surfaced as
   bare leaves ("returned a value of type Visibilities at output component [0]").

2. `via_image_from` called `transformer.visibilities_from(image=image)` WITHOUT
   `xp=xp`. That is the whole of the second fix — one line.

#422 characterised (2) as `TransformerNUFFT._forward_native` hard-converting to
NumPy and needing restructuring. That was wrong, and transformer.py is
deliberately untouched here: `_forward_native` already has a complete, jittable
JAX branch (lax.scan + dynamic_slice). The reported failure at transformer.py:660
was in the *NumPy* branch — the traceback pointed at the symptom while the cause
was one frame up, in a caller that never threaded xp. Same root shape as site 1
of #421.

TransformerDFT had been passing by luck: its arithmetic flows through tracers, so
the missing xp never showed. Only NUFFT, which calls into _nufftax and then
converts, exposed it.

Verified:
  - DFT and NUFFT, numpy-eager vs jax-jit, agree to ~1e-11 on real AND imaginary
    parts (complex visibilities)
  - both NUFFT chunk branches under jit match NumPy to ~5e-13 — single-shot and
    the lax.scan path (chunk_size is a transformer arg the simulator never sets,
    so that branch was exercised on the transformer directly)
  - the returned Interferometer keeps uv_wavelengths, real_space_mask, noise map
  - NumPy path still ndarray-backed
  - test_autoarray 929, test_autogalaxy 1009, test_autolens 488 — all unchanged

TransformerNUFFTPyNUFFT remains out of scope: the legacy pynufft backend is not
JAX-traceable and is not expected to be.

Closes #422

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01CWqjHGXUut25TEB8octU8H

Author: 2 people

autoarray/dataset/interferometer/simulator.py

@@ -7,6 +7,58 @@
 from autoarray import exc
 from autoarray.dataset import preprocess
 
+_INTERFEROMETER_PYTREES_REGISTERED = False
+
+
+def _register_interferometer_pytrees() -> None:
+    """Register ``Interferometer`` so ``jax.jit(via_image_from)`` can flatten its return.
+
+    Counterpart of ``_register_imaging_pytrees`` in the imaging simulator (and of
+    ``AnalysisImaging._register_fit_imaging_pytrees``). Without it a jitted
+    interferometer call raises ``TypeError: ... returned a value of type
+    Interferometer, which is not a valid JAX type``.
+
+    ``data`` and ``noise_map`` are the only per-simulation dynamic values.
+    Everything else rides as aux: ``uv_wavelengths`` and ``real_space_mask`` are
+    the fixed observation geometry, ``transformer`` and ``grids`` are constants for
+    a given simulation, the over-sample sizes are static integer geometry, and the
+    remaining slots are ``None``.
+
+    Unlike registration for a jitted function's *arguments* — which must happen
+    before the first call, since JAX flattens arguments at trace time — this is in
+    time, because the return value is flattened only after the body has run.
+    Idempotent via the module-level flag.
+    """
+    global _INTERFEROMETER_PYTREES_REGISTERED
+    if _INTERFEROMETER_PYTREES_REGISTERED:
+        return
+
+    from autoarray.abstract_ndarray import register_instance_pytree
+    from autoarray.structures.visibilities import Visibilities
+
+    # The two dynamic children must themselves be pytrees, or they surface as
+    # bare leaves and JAX rejects the return value ("... returned a value of type
+    # Visibilities ... at output component [0]"). Unlike ``Array2D`` — which the
+    # imaging path gets auto-registered — these are not registered anywhere else.
+    register_instance_pytree(Visibilities)
+    register_instance_pytree(VisibilitiesNoiseMap)
+
+    register_instance_pytree(
+        Interferometer,
+        no_flatten=(
+            "uv_wavelengths",
+            "real_space_mask",
+            "transformer",
+            "grids",
+            "over_sample_size_lp",
+            "over_sample_size_pixelization",
+            "noise_covariance_matrix",
+            "sparse_operator",
+        ),
+    )
+
+    _INTERFEROMETER_PYTREES_REGISTERED = True
+
 
 class SimulatorInterferometer:
     def __init__(
@@ -115,11 +167,14 @@ def via_image_from(self, image, xp=None):
         if xp is None:
             xp = self._xp
 
+        if xp is not np:
+            _register_interferometer_pytrees()
+
         transformer = self.transformer_class(
             uv_wavelengths=self.uv_wavelengths, real_space_mask=image.mask
         )
 
-        visibilities = transformer.visibilities_from(image=image)
+        visibilities = transformer.visibilities_from(image=image, xp=xp)
 
         if self.noise_sigma is not None:
             visibilities = preprocess.data_with_complex_gaussian_noise_added(