fix: bound MultiStartProdigy vmap memory (interferometer release-leg OOM)

[Jammy2211]

Overview

The nightly Release Integrate leg has been OOMing on scripts/interferometer/start_here.py with a single ~86 GB JAX allocation. The cause is not the NUFFT/transformer: af.MultiStartProdigy.batch_size defaults to None, which evaluates all n_starts in one jax.vmap and materializes the whole batched jvp at once. Both workspaces' start_here.py adopted MultiStartProdigy(n_starts=48) on 2026-07-29 with no batch_size; the next nightly failed. autolens_workspace@d7385ff already patched the autolens side with batch_size=4 — autogalaxy_workspace is still unguarded, and nothing in the library stops the next user hitting the same wall.

Plan

• Phase 1 (@autogalaxy_workspace) — mirror batch_size=4 into scripts/interferometer/start_here.py, the one script still exposed. Closes the live nightly failure.
• Phase 2 (@PyAutoHands) — set JAX_TRACEBACK_FILTERING=off at the single choke point every script run passes through, so a filtered traceback never again costs a night of triage.
• Phase 3 (@PyAutoFit) — make the memory guard automatic and correct, so MultiStartProdigy warns with an actionable batch_size instead of dying inside XLA.

Library-first ordering for merge; phase 1 is independent and can land immediately.

Detailed implementation plan

Affected Repositories

Repo	Phase	Role
PyAutoFit	3	primary — the general library fix
autogalaxy_workspace	1	the one script still unguarded
PyAutoHands	2	the script-runner env choke point

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	claude/interferometer-oom-nufft-yyz98z	clean
./autogalaxy_workspace	claude/interferometer-oom-nufft-yyz98z	clean
./PyAutoHands	claude/interferometer-oom-nufft-yyz98z	clean
./PyAutoMind	claude/interferometer-oom-nufft-yyz98z	clean

Branch: claude/interferometer-oom-nufft-yyz98z (session-designated; supersedes the usual feature/<task-name> convention). Named for the original NUFFT hypothesis, which the diagnosis below supersedes — the branch name is kept as-is rather than churned.

Evidence the cause is the vmap batch, not the NUFFT

• Only start_here.py uses this search. Audit of every search = af.* constructor under scripts/interferometer/ in both workspaces: start_here.py is the only user of af.MultiStartProdigy; all 45 others use af.Nautilus. In the failing run every sibling on the same dataset and the same TransformerNUFFT passed — fit.py, modeling.py, likelihood_function.py, all seven features/pixelization/*.
• The allocation divides by n_starts. 85,898,814,480 / 48 = 1,789,558,635 bytes ≈ 1.79 GB of jvp per start.
• Regression window. Both start_here.py files switched to MultiStartProdigy(n_starts=48) on 2026-07-29 (agw 255aee4 19:42, alw fa31bc7 21:15). First failure is the very next nightly, 07-30. A workspace authoring change, not a library change.
• Already reproduced. autolens_workspace@d7385ff (07-31 23:23) verified under the release profile env: unpatched reproduces the OOM at the same materialization site; patched runs end-to-end, exit 0, <4 GB peak.
• PyAutoFit already documents it. search.py:88-102 — "None (default) evaluates all n_starts in a single jax.vmap … allocates the whole batched jvp at once, which for a memory-heavy likelihood … exhausts even an 80 GB GPU."

Implementation Steps

Phase 1 — autogalaxy_workspace

1. scripts/interferometer/start_here.py: add batch_size=4 to the af.MultiStartProdigy(...) call, mirroring autolens_workspace@d7385ff including its explanatory comment.
2. Verify under the release profile (config/build/profile_release.yaml), not just smoke — the release leg is where it fails.
3. Regenerate the notebook.

Phase 2 — PyAutoHands

1. autohands/env_config.py::build_env_for_script: set JAX_TRACEBACK_FILTERING=off on the base env after the managed-prefix scrub and before apply_profile, so a profile can still override it.
2. Confirm the env_config is None path (inherit parent unchanged) is either covered or explicitly documented as out of scope.
3. Unit test in the env-config resolver tests.

Phase 3 — PyAutoFit

1. autofit/non_linear/analysis/analysis.py::print_vram_use — profile jax.value_and_grad(fitness.call), not fitness.call. It currently measures the likelihood, which under-reports precisely the gradient allocation that OOMs.
2. autofit/non_linear/search/mle/multi_start_gradient/search.py — when batch_size is None, project the batched-jvp footprint against available memory before compiling; on over-budget either emit an actionable warning naming a safe batch_size, or auto-batch. Must stay numerically inert (the existing _chunk_slices tiling already is).
3. Tests for the projection and the message.

Key Files

• autogalaxy_workspace/scripts/interferometer/start_here.py — n_starts=48, no batch_size; in smoke_tests.txt
• PyAutoHands/autohands/env_config.py — build_env_for_script / apply_profile
• PyAutoFit/autofit/non_linear/search/mle/multi_start_gradient/search.py — batch_size handling, _chunk_slices
• PyAutoFit/autofit/non_linear/analysis/analysis.py:337 — print_vram_use

Not implicated (checked, do not chase)

PyAutoArray/autoarray/operators/transformer.py — TransformerNUFFT.chunk_size, transform_mapping_matrix, and the nufftax 0.6.x _patch_nufftax_batchers shim. The guarded sibling scripts exercise all of this on the same dataset without incident.

Original Prompt

Click to expand starting prompt

The corrected prompt lives at PyAutoMind/active/interferometer_release_leg_oom.md. Its diagnosis section is reproduced above; the file additionally retains the superseded 2026-07-31 / 2026-08-04 NUFFT hypothesis as an auditable historical record, and the original symptom report:

integrate / run_scripts (3.12, autogalaxy, interferometer)   FAIL
integrate / run_scripts (3.12, autolens,    interferometer)  FAIL
  scripts/interferometer/start_here.py ...  FAIL (19.5s / 28.2s)
  For simplicity, JAX has removed its internal frames from the traceback of the
  following exception. Set JAX_TRACEBACK_FILTERING=off to include these.