feat(autofit): MultiStartProdigy + optax.contrib + per-start vmapped state

[Jammy2211]

Overview

Phase 1 of the multi-start gradient v2 promotion (autolens_workspace_developer#101). Promotes learning-rate-free / optax.contrib rule support into af.AbstractMultiStartGradient by resolving rules from optax.contrib, moving optimizer state from a stacked batch to per-start jax.vmap init/update (so lr-free rules' global scalar estimates no longer couple across starts), guarding steps with optax.apply_if_finite, and adding af.MultiStartProdigy. The headline #101 result: on the MGE cell Prodigy is bit-identical to hand-tuned Adam (+31787.84) with no learning rate. The restart-on-death recovery layer is deferred to Phase 2 (built on this PR's vmapped state).

Plan

• Resolve the optax rule from optax, falling back to optax.contrib (unlocks prodigy, dadapt, mechanic, dog/dowg, momo, schedule-free).
• Refactor optimizer state: stacked (n_starts, ndim) → per-start jax.vmap(opt.init) / vmapped update. Elementwise rules (Adam/ADABelief/Lion) stay numerically identical; lr-free rules stop sharing one global scalar across starts.
• Wrap the rule in optax.apply_if_finite as an in-step NaN guard; forward the value= kwarg for rules that need the loss (momo). Bump the optax pin to >=0.2.5.
• Add af.MultiStartProdigy (learning-rate-free) and export it from autofit/__init__.py.
• Numpy-only unit tests (contrib rule resolution, Prodigy defaults, dict round-trip); JAX MGE validation (prodigy == adam, per-start state independence) in autofit_workspace_test.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autofit_workspace_test (JAX validation, follow-on within the same task)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean

Suggested branch: feature/multistart-contrib-vmapped-state (worktree conflict check: clear)

Implementation Steps

1. autofit/non_linear/search/mle/multi_start_gradient/search.py — add a rule-resolver helper: getattr(optax, m, None) or getattr(optax.contrib, m), raising a clear error if neither has it. Replace both getattr(optax, self.optax_method)(self.learning_rate) sites (resume branch ~L193, fresh branch ~L204).
2. Same file, _fit: replace stacked opt_state = optimizer.init(params) and optimizer.update(grads, opt_state, params) with jax.vmaped init and a vmapped update step; wrap the resolved rule in optax.apply_if_finite(rule, max_consecutive_errors=...). Branch the update on whether the rule consumes value= (momo) — mirror pix_lr_free.py's needs_value / step_update. Preserve the search_internal resume path (opt_state is now a vmapped pytree — verify optax.tree_utils.tree_get round-trips).
3. Handle the learning-rate-free case: rules like prodigy take no learning_rate positional — resolver/subclass must construct them without forcing an lr (a MultiStartProdigy._default_learning_rate = None + rule-build indirection).
4. Add class MultiStartProdigy(AbstractMultiStartGradient) with optax_method = "prodigy"; export in autofit/__init__.py next to the existing three.
5. pyproject.toml — bump jax = [..., "optax"] to "optax>=0.2.5".
6. test_autofit/non_linear/search/mle/test_multi_start_gradient.py — numpy-only: Prodigy defaults/optax_method, contrib rule resolution, dict round-trip for Prodigy. Keep JAX out of the library suite.
7. autofit_workspace_test/scripts/searches/ — add a MultiStartProdigy.py JAX validation mirroring MultiStartAdam.py; assert prodigy reaches the Adam optimum on the MGE cell and that per-start state is independent.

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py — rule resolver, vmapped state, apply_if_finite guard, MultiStartProdigy.
• autofit/__init__.py — export MultiStartProdigy.
• pyproject.toml — optax>=0.2.5.
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py — numpy unit tests.
• autofit_workspace_test/scripts/searches/MultiStartProdigy.py — JAX MGE validation.
• Reference (read-only): autolens_workspace_developer/searches_minimal/pix_lr_free.py, lr_free_findings.md.

Testing

python -m pytest test_autofit/non_linear/search/mle/ (numpy). JAX MGE parity in autofit_workspace_test. Existing Adam/ADABelief/Lion results must be unchanged (elementwise rules are vmap-invariant).

Original Prompt

Click to expand starting prompt

Multi-start gradient search v2 — optax.contrib rules + per-start vmapped state

Type: feature
Target: autofit
Repos:

• @PyAutoFit
  Difficulty: medium
  Autonomy: supervised
  Priority: normal
  Status: draft

Phase 1 of the multi-start gradient v2 promotion (#101). The restart-on-death
(resurrection) recovery layer is Phase 2
(multistart_resurrection_restart_on_death.md), which builds on the per-start
vmapped state introduced here.

Promote the optax.contrib / learning-rate-free rule support from
autolens_workspace_developer#101 into af.AbstractMultiStartGradient
(autofit/non_linear/search/mle/multi_start_gradient/search.py):

1. optax.contrib rule resolution. The current
   getattr(optax, self.optax_method) misses optax.contrib (prodigy et al.).
   Resolve the rule from optax then fall back to optax.contrib.
2. Per-start vmapped optimizer state. The current stacked-(n_starts, ndim)
   state init COUPLES the lr-free rules' global scalar estimates (prodigy /
   dadapt d, DoG max_dist, mechanic scale, momo Polyak) across starts.
   Init/update the optimizer per start via jax.vmap (elementwise rules — the
   existing Adam/ADABelief/Lion — are numerically unaffected; the Remove use of sym link from Emcee, Dynesty. #101 benchmark
   showed no per-eval cost). This is the load-bearing refactor Phase 2 also
   depends on.
3. optax.apply_if_finite in-step guard. Wrap the rule so a single
   non-finite step zeroes the update instead of NaN-poisoning (optax >= 0.2.5
   forwards momo's value= kwarg). Bump the optax pin in pyproject.toml to
   optax>=0.2.5. (Resurrection in Phase 2 is the recovery layer on top of
   this guard, which alone only latches at the cliff edge.)
4. af.MultiStartProdigy. Add the concrete search — on the MGE cell prodigy
   is bit-identical to hand-tuned adam (+31787.84) with NO learning rate, the
   headline Remove use of sym link from Emcee, Dynesty. #101 result. Export it from autofit/__init__.py alongside the
   existing MultiStart* classes.

Wiring reference: autolens_workspace_developer/searches_minimal/pix_lr_free.py
(the per-start jax.vmap(opt.init) / vmapped update, the needs_value branch
for momo's value=). Evidence: searches_minimal/lr_free_findings.md
(#101, Phase 1 — MGE cell table).

Unit tests numpy-only (config knobs, rule resolution incl. contrib, dict
round-trip, Prodigy defaults); JAX validation (prodigy == adam on the MGE cell,
vmapped-state independence) in autofit_workspace_test per the established split.
Existing Adam/ADABelief/Lion results must be unchanged.

[Jammy2211]

Library PR Created — Phase 1

PR: #1398 (pending-release)

Promotes optax.contrib rule resolution + per-start jax.vmap optimizer state + optax.apply_if_finite guard + af.MultiStartProdigy (learning-rate-free) into af.AbstractMultiStartGradient; optax pin → >=0.2.5.

Validation (local): numpy unit suite 7 passed · JAX end-to-end (Adam/Prodigy/Prodigy-batch4) recover the Gaussian truth basin, Adam == Prodigy optimum with no learning rate · vmapped ApplyIfFiniteState dill-resume round-trips.

API changes: additive & backward-compatible — new af.MultiStartProdigy, new optional max_consecutive_nan arg; existing MultiStartAdam/ADABelief/Lion numerically unchanged. No workspace migration required (existing MultiStart references keep working).

Next: /start_workspace — author the permanent autofit_workspace_test/scripts/searches/MultiStartProdigy.py JAX validation (option i). Phase 2 (restart-on-death resurrection) remains a filed draft, to be issued after this merges.

Merge stays with the human (Heart was RED for pre-existing cross-repo/release-staleness reasons unrelated to this change; PyAutoFit gate green).

[Jammy2211]

Workspace PR Created — both PRs open, awaiting merge

• Library: feat(autofit): MultiStartProdigy + optax.contrib + per-start vmapped state #1398 (pending-release)
• Workspace: test(searches): JAX validation for af.MultiStartProdigy autofit_workspace_test#58 (pending-release) — adds scripts/searches/MultiStartProdigy.py JAX validation (passes: centre=50.156, normalization=25.196, sigma=9.858, no learning rate).

Merge order (library-first gate): merge #1398 first, then #58.

Both held for the human merge decision — Heart is RED for pre-existing, unrelated reasons (PyAutoLens dirty checkout, the known 632p/10f workspace-validation baseline predating this change, hygiene/release-staleness); this additive change's own gates are green (PyAutoFit unit suite 7-pass, JAX truth-basin recovery, dill-resume round-trip, workspace script passes).

Phase 2 (restart-on-death resurrection) remains a filed draft — issue it after #1398 merges.

[Jammy2211]

Shipped — Phase 1 complete. Both PRs merged: library #1398 and workspace PyAutoLabs/autofit_workspace_test#58. af.MultiStartProdigy (learning-rate-free) + optax.contrib support + per-start vmapped state + apply_if_finite guard are on main. Closing; Phase 2 (restart-on-death resurrection) tracked separately.