MultiStartGradient: Python-loop batched value_and_grad + jitted broad-start filter

[Jammy2211]

Summary

Multi-band af.FactorGraphModel fits driven by af.MultiStartProdigy paid an intractable cold JAX compile (>1 h on CPU; the batched jax.lax.map scan is memory-explosive to compile and was OOM-killed even at batch 1) plus a cache-immune ~13 min eager broad-start filter in every process. This PR:

• replaces the in-XLA lax.map scan used when batch_size is set with a Python sweep over jit(jax.vmap(value_and_grad)) chunks — the ragged final chunk is padded so exactly one (batch_size, ndim)-shaped program is compiled per search (pure _chunk_slices helper, numpy-only unit tested);
• jits the broad-start filter's single-point objective once in _fit and reuses it across draws (it previously rebuilt an eager jax.value_and_grad per draw — cache-immune, ~13 min per multi-band process).

Production-path result (homogeneous 4-band jwst MGE FactorGraphModel, 16 starts, batch_size=4; 1-core WSL host + RTX 2060): cold fit intractable → 395 s CPU / 392 s GPU; warm 136 s / 199 s. best_fom bit-identical across every arm, backend and cache state. New GPU probe rows show the scan explosion is CPU-backend-specific and steady eval is at parity (0.740 vs 0.729 s per 16-start sweep), so no scan fallback is kept. Closes #1430's library leg; census rows + findings note land in the companion autolens_profiling PR.

API Changes

None — internal changes only. batch_size semantics are unchanged (numerically inert tiling; still the memory bound, now also the compile bound).
See full details below.

Test Plan

• python -m pytest test_autofit/ — 1592 passed, 1 skipped (includes new numpy-only _chunk_slices tests: exact division, ragged, n < batch, batch=1, invariants)
• Parity: batch_size=None vs 4 vs ragged 5 — bit-identical best_fom / best_params / per-start params on the 1D-Gaussian JAX cell
• autofit_workspace_test/scripts/searches/MultiStartProdigy.py passes unchanged (truth basin recovered)
• Multi-band e2e cold/warm on CPU and laptop GPU (numbers above)

Full API Changes (for automation & release notes)

Removed

• autofit.non_linear.search.mle.multi_start_gradient.search.jax_value_and_grad_single(fitness, vector) — module-level helper, never exported via af.*; superseded by a jitted single-point objective built inside _fit. No usages found in any workspace repo.

Changed Behaviour

• AbstractMultiStartGradient(batch_size=...) (all MultiStart* searches) — the chunked sweep now runs from a Python loop over jit(jax.vmap(value_and_grad)) chunks instead of a single jax.jit(lax.map(...)) program. Numerically inert (bit-identical results); bounds both memory and XLA compile per chunk. Broad-start filtering now evaluates draws through one jitted (persistently cached) objective instead of eager per-draw AD — finite/non-finite classification of draws is unchanged in all validation runs.

Generated by the PyAutoLabs agent workflow.