Lift the nufftax cap to 0.6.x — batched nufft2d2 vjp (nightly release blocker) #424

Problem

The nightly release's Stage 3 (release-fidelity workspace validation) has failed 4 of the last 5 nights (e.g. PyAutoHeart run 30516167217): autogalaxy and autolens scripts/interferometer/start_here.py crash with

nufftax/transforms/autodiff.py:676, in _nufft2d2_bwd
    n2, n1 = f.shape
ValueError: too many values to unpack (expected 2)

Trigger: workspace commits on 2026-07-29 19:42 (autolens_workspace d5c9802d, autogalaxy_workspace 255aee4) made MultiStartProdigy the default search in interferometer/start_here.py — the first time CI differentiates through TransformerNUFFT. The library diff since the last green rehearsal is innocent (transformer.py unchanged).

Root cause: nufftax 0.4.0's _nufft2d2_bwd custom-vjp assumes 2-D f, but autoarray/operators/transformer.py transform_mapping_matrix feeds a batched stack. Reproduced with a minimal batched jax.value_and_grad over nufft2d2.

Fix: nufftax 0.6.x supports batched input natively (f.ndim == 3) with primitive-based autodiff — the same repro passes, and the forward transform is bit-identical to 0.4.0 at autoarray's eps=1e-12 (2-D + batched, x64), so no pinned-likelihood baselines move. The old <0.5.0 cap existed only because nufftax ≥0.5 needed Python ≥3.12 — the ecosystem floor is now 3.12.

Plan

High-level

• Lift nufftax>=0.4.0,<0.5.0 → nufftax>=0.6.1,<0.7.0 everywhere it is pinned.
• Verify end-to-end by running the previously-failing script locally on 0.6.1.
• No source or test changes (library unit tests are numpy-only; the release-mode workspace validation is this path's regression net).

Detailed

1. PyAutoArray PR — pyproject.toml: main dependency + dev extra, both to nufftax>=0.6.1,<0.7.0. Full unit suite before push.
2. PyAutoHeart PR — .github/workflows/workspace-validation.yml: the three pip install \"nufftax>=0.4.0,<0.5.0\" lines → the new range; update the stale "jax 0.10.x + nufftax 0.4.x is the tested combo" comment.
3. Evidence run: autogalaxy_workspace/scripts/interferometer/start_here.py under PYAUTO_TEST_MODE=1 with nufftax 0.6.1 — must get past MultiStartProdigy._broad_starts (the CI crash point at ~20 s).

Branch: feature/nufftax-cap-lift · Worktree: ~/Code/PyAutoLabs-wt/nufftax-cap-lift/ (PyAutoArray + PyAutoHeart)
Prompt: PyAutoMind/active/nufftax_batched_nufft2d2_vjp_release_blocker.md