feat: warn before an unbatched MultiStartProdigy jvp exhausts memory

[Jammy2211]

Phase 3 of 3 for #1452.

What and why

With batch_size=None, MultiStartProdigy vmaps every start into one value_and_grad. The only signal that this will not fit is XLA's RESOURCE_EXHAUSTED: Out of memory allocating N bytes, raised from inside the compiled program with nothing pointing at batch_size — the one knob that fixes it. That cost two nightly release runs in 2026-07, after which two workspace scripts were hand-patched with magic numbers.

Before the full vmap is compiled, this now probes the batched jvp at widths 1 and 2 via XLA's own memory analysis, splits the result into fixed and per-start components, and warns naming a batch_size that would fit.

Two probe points, not one, on purpose. Workspace guidance claimed VRAM "does not scale with batch size for the persistent buffers, so if it fits at batch_size=1 you can push it up (e.g. to 50)" — true only when the per-start slope is small next to the fixed cost. For an interferometer jvp it is not: the incident measured ~1.79 GB per start, so 48 starts wanted ~86 GB while the single-start probe fitted comfortably. A one-probe model reproduces the bad advice. (That prose is corrected in PyAutoLabs/autolens_workspace#… on the matching branch.)

Analysis.print_vram_use gains a gradient flag and now names what it measured. It profiled fitness.call — the likelihood, not the jvp — so its number could never bound a gradient search's footprint.

Design choices worth reviewing

• Warns only; never auto-applies. batch_size is documented numerically inert, so auto-batching would be safe in principle — but silently changing the execution shape of every existing run on the strength of a projection isn't a trade to make unattended. Escalating to auto-batch is a reasonable follow-up once the probe has run on real hardware.
• Fails open. Any exception in the projection is swallowed. A memory estimate must never be why a fit doesn't start.
• Public API: Analysis.batched_memory_bytes is new; print_vram_use gains an optional gradient=False (existing calls unchanged, though their output string now names the mode).

Known limitation — please read

memory_analysis reports 0 bytes on a CPU-only JAX build, which is where the release harness runs. The projection is skipped there rather than guessing, so this does not fix the CI OOM — it helps GPU users, and the CPU path relies on the now-unfiltered traceback from PyAutoLabs/PyAutoHands (phase 2). Documented at the call site. Making the CPU path measurable needs a live run to validate and is deliberately not attempted here.

Testing

The projection arithmetic is pure Python and unit-tested in test_multi_start_gradient.py, matching that suite's existing boundary ("plumbing is pure NumPy and tested here… keeping JAX out of the library unit suite"). Cases cover recovery of the incident's exact per-start cost, non-monotonic/identical probes (no negative slope), and six degenerate budget combinations.

The JAX probe itself has not been executed. It was authored in an environment with no jax, autofit or pytest installed, so pytest test_autofit/ has not been run and batched_memory_bytes has never called into XLA. Please run the suite and one real MultiStartProdigy fit before merging.

---

Generated by Claude Code