Problem
AbstractMultiStartGradient (PyAutoFit#1369/#1370) vmaps every start at once:
batched_value_and_grad = jax.jit(jax.vmap(jax.value_and_grad(fitness.call)))
For memory-heavy likelihoods this OOMs. Measured on an A100 80GB with a pixelized (kernel-CDF mesh) lens likelihood at n_starts=16, float64: the vmapped jvp fusion is f64[16, 15361, 2, 31, 512] = 58.13 GiB → RESOURCE_EXHAUSTED. There is currently no knob to reduce it — the only workarounds are dropping to float32 (a science compromise) or cutting n_starts (defeats multi-start).
Prior A100 work recorded the same class of failure ('OOMs on pix/delaunay via vmap').
Fix
Add batch_size: Optional[int] = None to AbstractMultiStartGradient (inherited by MultiStartAdam / MultiStartADABelief / MultiStartLion): chunk the vmapped value_and_grad over starts in batches of batch_size, concatenating results. None = current all-at-once behaviour (no regression). The optax update still runs on the full stacked array — only the memory-dominant jvp is chunked.
Precedent: af.Nautilus already exposes n_batch, and autolens_profiling maintains a per-(dataset, model, instrument) vmap batch_size table for the A100 (vram/config.py; e.g. imaging/delaunay at HST scale needs batch=16, not the default 100). This makes that table consumable by the gradient searches too.
Test plan
- numpy-only unit tests: batch_size plumbing/defaults, dict round-trip, and that chunked == unchunked on a simple objective.
- JAX end-to-end validation in autofit_workspace_test.
Blocks the pixelized multi-start experiment (autolens_workspace_developer#100).
Problem
AbstractMultiStartGradient(PyAutoFit#1369/#1370) vmaps every start at once:For memory-heavy likelihoods this OOMs. Measured on an A100 80GB with a pixelized (kernel-CDF mesh) lens likelihood at
n_starts=16, float64: the vmapped jvp fusion isf64[16, 15361, 2, 31, 512]= 58.13 GiB →RESOURCE_EXHAUSTED. There is currently no knob to reduce it — the only workarounds are dropping to float32 (a science compromise) or cutting n_starts (defeats multi-start).Prior A100 work recorded the same class of failure ('OOMs on pix/delaunay via vmap').
Fix
Add
batch_size: Optional[int] = NonetoAbstractMultiStartGradient(inherited byMultiStartAdam/MultiStartADABelief/MultiStartLion): chunk the vmappedvalue_and_gradover starts in batches ofbatch_size, concatenating results.None= current all-at-once behaviour (no regression). The optax update still runs on the full stacked array — only the memory-dominant jvp is chunked.Precedent:
af.Nautilusalready exposesn_batch, and autolens_profiling maintains a per-(dataset, model, instrument) vmap batch_size table for the A100 (vram/config.py; e.g. imaging/delaunay at HST scale needs batch=16, not the default 100). This makes that table consumable by the gradient searches too.Test plan
Blocks the pixelized multi-start experiment (autolens_workspace_developer#100).