MultiStartGradient: Python-loop batching fixes multi-band compile

[Jammy2211]

Overview

Multi-band af.FactorGraphModel fits driven by af.MultiStartProdigy (the autolens_assistant COSMOS-Web Ring case: F277W+F444W JWST fitted simultaneously) pay a cold JAX compile exceeding one hour on CPU even with every shipped mitigation live (persistent cache, autotune off, constant-folding disabled). The compile census (autolens_profiling#93) localised the cause: the in-XLA jax.lax.map scan used by the batch_size path is compile-intractable (and memory-explosive to compile — OOM-killed at ~6 GB RSS even at batch 1) precisely when its body is a multi-band FactorGraphModel fusion, while the same width via a Python loop over jit(vmap(value_and_grad)) chunks ("pyloop") compiled in 166 s. Single-band compile is benign either way. This task productizes the validated pyloop lever in AbstractMultiStartGradient, speeding up cold compile on CPU and GPU alike.

Plan

• Replace the in-XLA lax.map batched path in AbstractMultiStartGradient._fit with Python-loop batching: compile jax.jit(jax.vmap(value_and_grad)) over one batch_size-shaped chunk and sweep chunks in Python.
• Pad the ragged final chunk (broad-start collection may return fewer than n_starts rows) so the whole search pays exactly one chunk compile; discard padded outputs.
• Keep batch_size semantics: still the memory bound, now also the compile/throughput tunable (small on CPU, wider on GPU). batch_size=None (single full vmap) unchanged.
• Numpy-only unit tests for the chunk/pad helper; JAX parity + benchmarks via workspace-test scripts and the autolens_profiling jax_compile probe.
• Benchmark cold compile + steady eval, old vs new path, multi-band datacube cells on CPU and laptop GPU; record in autolens_profiling notes + README.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autolens_profiling (benchmark validation + durable findings notes)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autolens_profiling	main	clean (untracked benign dataset/imaging/jwst_lw/)

Suggested branch: feature/multiband-pyloop-batching

Note: worktree_check_conflict flagged PyAutoFit as claimed by task python-312-floor, but that Codex worktree's branch is 0 commits ahead of origin/main (PR #1429 merged) — stale claim, not a real conflict.

Sizing note (Brain override)

The Brain FeatureDecision scored this too-large (score 13) off repo count and proposed a generic 4-phase split. Overridden to a single implement+benchmark task: the lever is already validated (autolens_profiling#93 + the deeper-dig prompt's sub-investigation B), the library change is one method plus tests, and the benchmark closes the loop. Precedent: the #93 measure→fix override.

Implementation Steps

1. autofit/non_linear/search/mle/multi_start_gradient/search.py — in _fit, delete the @jax.jit + jax.lax.map(_value_and_grad, params, batch_size=...) construction (currently lines 252–259). When batch_size is set, keep the single _vmapped = jax.jit(jax.vmap(_value_and_grad)) and define batched_value_and_grad(params) as a Python sweep: iterate params in batch_size chunks; pad the ragged final chunk to batch_size (repeat last row) and trim the padded outputs; concatenate per-chunk (foms, grads). The jit cache keys on the (batch_size, ndim) shape → exactly one chunk compile per search.
2. Factor the chunk/pad bookkeeping into a pure helper _chunk_slices(n_rows, batch_size) -> [(lo, hi, pad)] so it is numpy-only unit-testable (library unit tests never import JAX).
3. Update the batch_size docstring: drop the lax.map description; state the Python-loop tiling, unchanged memory-bound semantics, and the compile rationale (in-XLA scan over a multi-band FactorGraphModel fusion is compile-intractable; the Python loop compiles one chunk).
4. step_update (optax, elementwise) stays as-is — not the compile driver.
5. test_autofit/non_linear/search/mle/test_multi_start_gradient.py — add numpy-only tests for _chunk_slices (exact division, ragged, n < batch, batch=1); update test__batch_size_is_carried_to_every_rule's docstring, which cites the jax.lax.map property.

Validation & benchmarks

1. python -m pytest test_autofit/ stays green (numpy-only).
2. Parity: batch_size=4 vs batch_size=None foms/grads identical on a small JAX problem via autofit_workspace_test/scripts/searches/MultiStartProdigy.py.
3. Headline A/B: multi-band FactorGraphModel + MultiStartProdigy cold compile, old lax.map vs new pyloop, CPU and laptop GPU (fresh JAX_COMPILATION_CACHE_DIR per arm; XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 on laptop GPU; idle machine — compile is host-CPU-bound). Instruments: autolens_profiling/scripts/misc/jax_compile/probe.py datacube cells (datacube_img / datacube_img_hetero, laxmap_vag vs pyloop_vag) plus one end-to-end multi-band fit timing (cold + warm).
4. Record results in autolens_profiling/results/notes/ + jax_compile/README.md multi-band section; keep ruff clean.

Fallback (decision on evidence, not implemented up front)

If the GPU steady-eval benchmark shows a real throughput regression from host dispatch (~12 dispatches/step at defaults; expected negligible vs likelihood cost), reintroduce the scan behind an explicit opt-in rather than as the default.

Out of scope (follow-ups)

• A100 census rows (RAL confirmatory tier; A100s saturated during GridSearch pickling #93 — re-run census_gpu.sbatch later).
• The padding/shape-canonicalization helper for heterogeneous bands (prompt item 3; secondary lever — the transform, not heterogeneity, is the dominant driver).
• Re-opening the single-band Uniform prior should not set limits on parameter in later phases. #71→Output mask pickle for phase #77 "settings suffice" verdict.

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py — the batched value_and_grad construction in _fit.
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py — numpy-only knob/helper tests.
• autolens_profiling/scripts/misc/jax_compile/probe.py — validated pyloop_vag pattern + benchmark instrument.
• autolens_profiling/jax_compile/README.md (via scripts/misc/jax_compile/README.md) — multi-band findings section to update.

Original Prompt

Click to expand starting prompt

Multi-band FactorGraphModel value_and_grad compile — deeper dig + productize the fix

Type: research
Target: PyAutoFit
Repos:

• PyAutoFit
• autolens_profiling
  Difficulty: large
  Autonomy: supervised
  Priority: normal
  Status: formalised

Follow-up to the multi-band compile investigation. Experiment A is already done
(autolens_profiling branch research/multiband-compile-ab; census in
autolens_profiling/jax_compile/README.md "Multi-band … heterogeneous-shape cliff"
section). Established, on local CPU MGE vag over a 4-band af.FactorGraphModel:

• Homogeneous 4-band (1 distinct shape) == single-band: cold compile 120s
  (≈ the 117s single-band figure). XLA fuses identical-shape factors into one
  shared kernel; the factor graph adds no compile cost when band shapes match.
• Heterogeneous (2 distinct shapes) = 704s cold = 5.9× the control, superlinear
  in the number of distinct shapes (2 shapes → ~6×, not 2×). Trace and steady
  eval unchanged → pure XLA fusion-compilation effect. This reproduces the real

  1h cold compile of an N-distinct-shape multi-wavelength fit.

• Persistent cache rescues both arms (multi-band caching now certified): warm
  compile 2–7s. The cliff is a one-time first-compile (cache-miss) cost per graph
  structure.

Reproduce via the probe cells added in Experiment A: python jax_compile/probe.py --dataset-class datacube_img[_hetero] --model-type mge --transforms vag --cache-dir <fresh> (run twice cold/warm). Two distinct shapes come from
jwst (0.03") + jwst_lw (0.06") channels.

Reproduction-context follow-up (2026-07-21) — what actually causes the >1h.
Heterogeneity ALONE does not reach the reported >1h: hetero vag topped out at
~12 min. Profiling the MultiStartProdigy transform (lax.map/vmap over
starts) on the SAME (single-core, nproc=1) host, homogeneous 4-band graph,
located the dominant driver:

• vag (start-width 1): 120 s compile
• vmap_vag (width 2): 209 s compile
• laxmap_vag (MultiStartProdigy default: batch_size 4 / n_batch 16):
  did not finish compiling in 55 min (killed).

So compile scales steeply with multi-start width, and the full production
transform is intractable to compile cold on one core even before heterogeneity —
which then multiplies it by 5.9×. Driver order: transform × core count >>
heterogeneity >> the factor graph itself. The nproc=1 dev host is worst-case
for XLA compile (host-CPU-bound); HPC/A100 numbers will be far lower. Biggest
shipped lever is the persistent cache (amortizes the whole cold cost). This is
recorded in the jax_compile/README.md multi-band section (findings 2 + 4).

This task is the deeper dig + the fix, as one task:

1. Sub-investigation B — the source lever (now concrete + validated). The
   dominant compile cost is the in-XLA lax.map scan in MultiStartProdigy, NOT
   the batching itself: at vmap width 1, pyloop_vag (Python loop over vmap
   chunks, batching hoisted out of XLA) compiled in 166 s vs laxmap_vag (in-XLA
   lax.map scan) intractable (>30 min). Implement Python-loop multi-start
   batching in MultiStartProdigy (the pyloop pattern in
   autolens_profiling/jax_compile/probe.py): iterate starts in Python over small
   vmap chunks so no scan is compiled; batch_size = vmap width becomes a
   compile/throughput tunable (small CPU, wider GPU). Benchmark cold compile +
   per-eval runtime, especially on GPU/A100 vs the current lax.map path —
   pyloop trades on-device loop fusion for host dispatch overhead and loses
   composability (can't jit/scan over the whole search), so the win is clear on
   CPU/limited-RAM but is a compile-vs-throughput trade on A100 (pick chunk width
   per device). Clean re-confirm DONE (2026-07-21): laxmap bs=1 compile was
   OOM-killed (~6 GB RSS) even from a 10 GB-free start → the lax.map scan path is
   memory-explosive to compile here, confound resolved in favour of the lever.
   Secondary: whether a per-factor jit boundary inside
   af.FactorGraphModel.log_likelihood_function additionally bounds the
   heterogeneity multiplier (N×single-band + linear combine) without padding.
2. Sub-investigation C — multi-core / A100 rows + verdict, including the
   transform. Add A100 (and/or multi-core CPU) multi-band rows for BOTH vag
   AND the production laxmap_vag transform (datacube_img /
   datacube_img_hetero) to the jax_compile census — the transform is the
   dominant driver (see reproduction context) and the single-core figure is
   worst-case, so the central open question is whether the full MultiStartProdigy
   compile is tractable on an HPC host (single-band A100 vag cold was ~28s vs
   229s CPU). Quantify the transform-width scaling and heterogeneity multiplier
   under real core counts, then ship the final N-band compile verdict.
3. Productize the immediate user workaround — pad short-wavelength bands to a
   common grid so all factors share one fused kernel — as a documented recipe
   and/or helper for N-band gradient fits.

Deliverable: multi-band A100 census rows + either a productizable compile-reduction
lever (per-factor jit or the padding helper) or a documented "GPU-only /
same-shape-bands" verdict. Out of scope: re-opening the #71 single-band
settings verdict.

[Jammy2211]

Progress: pyloop batching implemented in AbstractMultiStartGradient._fit (worktree feature/multiband-pyloop-batching).

• Full PyAutoFit unit suite: 1592 passed, 1 skipped (new numpy-only _chunk_slices tests included).
• Parity: batch_size=None vs 4 vs ragged 5 on the 1D-Gaussian JAX cell — bit-identical best_fom / best_params / all per-start params, same step counts.
• Canonical autofit_workspace_test/scripts/searches/MultiStartProdigy.py control run passes unchanged.

Now benchmarking the production path on the 4-band datacube_img cell (CPU cold/warm e2e, then laptop-GPU rows incl. a time-capped laxmap_vag comparison).

[Jammy2211]

Shipped to PR-open — merge left to human review.

• Library PR: MultiStartGradient: Python-loop batched value_and_grad + jitted broad-start filter #1431 (pending-release verified)
• Census/findings PR: Multi-band pyloop productized: findings note + GPU census rows (PyAutoFit#1430) autolens_profiling#95 (pending-release verified)

Production-path benchmark (4-band jwst MGE FactorGraphModel, MultiStartProdigy 16 starts / batch_size=4; 1-core WSL host, RTX 2060):

arm	cold fit	warm fit
CPU, old lax.map path	intractable (>55 min, OOM-killed)	—
CPU, shipped (pyloop + jitted broad-starts)	395.5 s	136.2 s
GPU, shipped	392.3 s	198.5 s

best_fom bit-identical across all arms/backends/cache states. Second fix folded in: the broad-start filter was evaluating an eager per-draw value_and_grad (cache-immune, ~13 min/process on multi-band) — now one jitted, persistently-cached objective. New census finding: the scan explosion is CPU-backend-specific (GPU compiles laxmap_vag in ~122 s at steady-eval parity 0.740 vs 0.729 s/sweep) → the PR-noted opt-in scan fallback is not needed.

Validation: PyAutoFit suite 1592 passed / 1 skipped · parity batch_size None/4/ragged-5 bit-identical · targeted smoke: MultiStartProdigy.py, MultiStartAdam.py, MultiStartResurrect.py, multi_start_gradient_auto_convergence.py all pass.

Heart at ship time (surfaced verbatim, pre-existing/unrelated): verdict stale — "test run status unknown (no report.json)"; "release validation stale: source moved since rehearsal (PyAutoFit, PyAutoArray, PyAutoGalaxy, PyAutoLens)". Vitals otherwise green (install verify + release validation last runs passed 2026-07-29).

Follow-ups (kept in the active/ prompt): A100/multi-core census rows (census_gpu.sbatch re-run when A100s free); datacube_img_hetero GPU rows; padding/shape-canonicalization helper; per-factor jit boundary (secondary).

[Jammy2211]

Shipped: PyAutoFit#1431 merged (f5bdc7c6), autolens_profiling#95 merged (014afbbf). Cold multi-band MultiStartProdigy fit: intractable → ~6.5 min (CPU and laptop GPU); warm ~2–3 min; bit-identical numerics. Follow-ups (A100/multi-core + hetero-GPU census rows; padding-helper / per-factor-jit reassessment) re-filed as a Mind draft prompt.