feat(search): add batch_size to the multi-start gradient searches

[Jammy2211]

Summary

The multi-start gradient searches (#1369/#1370) vmap every start's value_and_grad at once. For memory-heavy likelihoods that OOMs, with no knob to reduce it.

Measured on an A100 80GB, pixelized (kernel-CDF mesh) lens likelihood, n_starts=16, float64:

f64[16, 15361, 2, 31, 512] = 58.13 GiB  ->  RESOURCE_EXHAUSTED

The only workarounds were float32 (a science compromise) or fewer starts (defeats multi-start).

Fix

batch_size: Optional[int] = None on AbstractMultiStartGradient (inherited by MultiStartAdam / MultiStartADABelief / MultiStartLion).

• None (default) → unchanged: one jax.vmap over all starts. No regression.
• Set → tiling handed to jax.lax.map(..., batch_size=), JAX's own chunked map: vmaps within a chunk, scans across chunks, and handles a ragged final chunk without a second compile.
• lax.map is used only when batching is requested — with batch_size=None it degrades to a sequential scan, which would destroy the default path's parallelism.

Only the memory-dominant gradient evaluation is tiled; the optax update still runs over all starts.

Not the same thing as Nautilus's n_batch

af.Nautilus's n_batch is Nautilus's own algorithmic knob (how many points it proposes per iteration) that autofit forwards. Here n_starts is the algorithm; batch_size is purely implementation-level tiling and is numerically inert.

(For the record: optax offers nothing for this — MultiSteps/apply_every are gradient accumulation over steps for a single parameter set fed many data batches. Our starts are independent parameter sets; accumulating across them would change the algorithm.)

Test Plan

• test_autofit/ full suite: 1480 passed, 1 skipped (no regression)
• Numerically inert through the real _fit path: best_fom/best_params identical for batch_size ∈ {None, 1, 4, 14, 100} — incl. ragged (14 % 4 == 2) and batch > n_starts
• numpy-only plumbing tests (defaults None, carried by all three rules)
• JAX equivalence assertion in autofit_workspace_test (follow-up, library-first)

Unblocks the pixelized multi-start experiment (autolens_workspace_developer#100), where the vram batch table can now be consumed by the gradient searches as it already is by Nautilus.

Generated by the PyAutoLabs agent workflow.