fix(config): release profile disables JAX, breaking MultiStartAdam validation shard

[Jammy2211]

Summary

nightly-release has failed five nights running (2026-07-11 → 2026-07-15). Stage 3 (release-fidelity integration) dispatches PyAutoHeart workspace-validation.yml at mode=release, and exactly one shard fails — run_scripts (3.12, autofit_test, searches):

scripts/searches/MultiStartAdam.py ...   FAIL (1.7s)
ValueError: MultiStartAdam is a JAX-native gradient search and requires a JAX-traceable
Analysis (e.g. `AnalysisImaging(..., use_jax=True)`). The supplied analysis is not running
on the JAX backend.

Latest: nightly 29390030719 → validation 29390440297 (job 87272559383). Every other shard in that run is green.

This is not the shifting perf-flake tail from the release-tail epic — it is a fresh, deterministic regression from the multi-start gradient search promotion (PyAutoFit #1370 + workspace follow-ups, merged 2026-07-14). This repo's HEAD is literally 0a81c45 test(searches): validate MultiStartAdam recovers truth basin under JAX (#43).

Root cause

MultiStartAdam._fit hard-guards on the JAX backend (PyAutoFit .../multi_start_gradient/search.py:120-126):

if not getattr(analysis, "_use_jax", False):
    raise ValueError(...)

config/build/env_vars_release.yaml — the mode=release profile, not env_vars.yaml which is the per-PR smoke profile — pins:

• PYAUTO_TEST_MODE: "0" → real searches, so AbstractSearch.fit does not take the mode >= 2 bypass (abstract_search.py:683-689) and _fit really runs.
• PYAUTO_DISABLE_JAX: "1" → forces use_jax=False, so the script's af.ex.Analysis(..., use_jax=True) is silently flipped and the guard fires.

The profile re-enables JAX for exactly two patterns (jax_assertions/, searches/BlackJAXNUTS). searches/MultiStartAdam was never added.

Why sibling scripts don't fail

• autofit_workspace/scripts/searches/mle.py uses af.MultiStartAdam with use_jax=True, but that workspace's env_vars_release.yaml pins PYAUTO_DISABLE_JAX: "0" — JAX on, passes.
• autolens_workspace / autogalaxy_workspace guides/modeling/searches.py construct the search but never call .fit() — config-only, guard never reached.

Why the per-PR gate couldn't catch it

Under smoke (env_vars.yaml), PYAUTO_TEST_MODE: "2" bypasses the sampler before _fit, so MultiStartAdam.py passes without ever executing its JAX path. The gate was structurally blind to this.

High-level plan

• Add a release-profile override so MultiStartAdam runs with JAX enabled — the blocking fix.
• Do the same for Dynesty_jax / Nautilus_jax, which declare use_jax=True but were also missing from the overrides and so were passing vacuously on the NumPy path.
• Make MultiStartAdam really execute under the smoke profile, so the per-PR gate can catch this class of regression at PR time instead of on night five.
• Verify every affected script actually passes with JAX on before shipping.

Detailed plan

1. config/build/env_vars_release.yaml — three set: overrides (the profile's header documents that overrides here only ever set:, never unset:, a var defaults pins):

  - pattern: "searches/MultiStartAdam"
    set: { PYAUTO_DISABLE_JAX: "0" }
  - pattern: "searches/Dynesty_jax"
    set: { PYAUTO_DISABLE_JAX: "0" }
  - pattern: "searches/Nautilus_jax"
    set: { PYAUTO_DISABLE_JAX: "0" }

2. config/build/env_vars.yaml (smoke) — mirror the searches/BlackJAXNUTS precedent:

  - pattern: "searches/MultiStartAdam"
    unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX]

Scoped to MultiStartAdam only; the nested samplers stay bypassed under smoke on cost grounds.

Trade-off considered: leaving the smoke profile alone was the leaner option, but it preserves the exact blind spot that let this reach night five. Human's explicit call: it should not stay bypassed if running it is what ensures the bug is caught. Cost is ~6s per PR.

Verification (all run on feature/multistart-adam-release-jax)

Scripts run with the release-profile env (PYAUTO_TEST_MODE=0 PYAUTO_DISABLE_JAX=0 JAX_ENABLE_X64=True):

Script	Result
MultiStartAdam.py	exit 0 — recovered centre=50.156, normalization=25.196, sigma=9.858 (truth 50/25/10)
Dynesty_jax.py	exit 0 (~10s)
Nautilus_jax.py	exit 0 (~25s)

MultiStartAdam.py under the smoke-resolved env (both vars unset): exit 0, same recovery, ~6s.

No bug was hiding behind the two vacuous passes, so all three overrides are safe to ship.

Resolver check — drove autobuild.env_config.build_env_for_script over the real configs for every scripts/searches/*.py, guarding against the substring/.py-anchor over-match trap:

===== env_vars.yaml (smoke) =====        ===== env_vars_release.yaml =====
script              TEST_MODE  JAX       script              TEST_MODE  JAX
BlackJAXNUTS.py       <unset>  <unset>   BlackJAXNUTS.py             0    0
MultiStartAdam.py     <unset>  <unset>   MultiStartAdam.py           0    0
Dynesty_jax.py              2        1   Dynesty_jax.py              0    0
Nautilus_jax.py             2        1   Nautilus_jax.py             0    0
Nautilus.py                 2        1   Nautilus.py                 0    1
DynestyStatic.py            2        1   DynestyStatic.py            0    1
DynestyDynamic.py           2        1   DynestyDynamic.py           0    1
Emcee.py / LBFGS.py         2        1   Emcee.py / LBFGS.py         0    1
Zeus.py                     2        1   Zeus.py                     0    1

Nautilus.py and DynestyStatic/Dynamic.py correctly stay at DISABLE_JAX=1 — the patterns did not over-match.

Acceptance

• MultiStartAdam.py passes under the release profile and recovers the truth basin
• Dynesty_jax.py / Nautilus_jax.py verified green with JAX enabled
• MultiStartAdam.py really executes (not bypassed) under smoke
• autofit_test/searches shard green; nightly-release reaches Stage 3 clean

---

Prompt: PyAutoMind/active/bug/autofit_workspace_test/multistart_adam_release_profile_disables_jax.md
Branch: feature/multistart-adam-release-jax
Autonomy: supervised (bug cap) — plan written here per the autonomy contract; ship sign-off parks for human review; merge stays human.

[Jammy2211]

⏸ Parked: ship sign-off

Implementation and verification are done. bug caps autonomy at supervised, so this is checkpoint 2 (ship PR sign-off) — I need a human go-ahead before commit/push/PR. Merge stays human regardless.

Scripts Changed

No script or library source changed — config only, 2 files, +23 lines (comments included):

File	Change
config/build/env_vars_release.yaml	+3 set: { PYAUTO_DISABLE_JAX: "0" } overrides — searches/MultiStartAdam, searches/Dynesty_jax, searches/Nautilus_jax
config/build/env_vars.yaml	+1 unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX] override — searches/MultiStartAdam

API Changes

None. No public API, no library source, no script logic touched.

Ship-gate legs

Leg	Status
1. Tests	n/a — autofit_workspace_test has no library test dir; the scripts are the tests, and all four affected ones were run directly (below).
2. Smoke	effectively run — the changed surface is the searches shard itself; every scripts/searches/*.py env resolution was verified against the real configs, and the 3 JAX scripts executed end-to-end.
3. Review	self-reviewed; diff is 23 lines of config + rationale comments.
4. Heart	YELLOW score 55, reasons: workspace-validation not passing (this bug), 58 stale parked scripts, install verification not run, release validation stale. Not acknowledged at launch, so this leg parks the run by contract. Notably the first reason is what this PR fixes.

Evidence

All four scripts pass with JAX enabled (release-profile env):

• MultiStartAdam.py → exit 0, recovered centre=50.156, normalization=25.196, sigma=9.858 (truth 50/25/10)
• Dynesty_jax.py → exit 0 (~10s)
• Nautilus_jax.py → exit 0 (~25s)
• MultiStartAdam.py under smoke-resolved env (both vars unset) → exit 0, ~6s

The two vacuous scripts had no bug hiding behind them, so all three release overrides are safe.

The question

Approve the ship? On approval I commit, push feature/multistart-adam-release-jax, and open a pending-release PR against autofit_workspace_test:main.

• Approve → PR opened; merge + close still yours; the next nightly should reach Stage 3 clean.
• Reject / amend → say what to change; branch stays local, nothing pushed.

One judgement worth a second look: the smoke-profile change (env_vars.yaml) adds ~6s of real JAX run to every PR. It was your explicit call — "it should not opt out if it ensures its caught" — and it is the only leg that would have caught this before night five. It is also the one line here that is a preference rather than a fix, so it is the easiest thing to drop if you would rather keep the per-PR gate cheap.

[Jammy2211]

⚠️ Correction — the smoke-profile rationale in the PR body and commit message is wrong

I claimed the env_vars.yaml change makes the per-PR gate catch this class of bug. That is not true, and I should correct it before merge rather than let it stand in main's history.

What I got wrong. There are two distinct things both loosely called "smoke":

1. The per-PR gate — .github/workflows/smoke_tests.yml, which runs the curated smoke_tests.txt subset (11 scripts: Emcee, DynestyStatic, Nautilus, features/*, database/*, …). Neither MultiStartAdam nor BlackJAXNUTS is in that list. This gate never runs either script, so the env_vars.yaml override does not affect it at all.
2. The smoke env profile — config/build/env_vars.yaml, consumed by PyAutoHeart's workspace-validation.yml at mode=smoke, which runs every script (weekly, Mondays 03:00 UTC, plus manual dispatch).

The env_vars.yaml change lands on (2), not (1).

What the change is actually worth. Still real, just on a different cadence than I said. Before it, the weekly full-validation run executed MultiStartAdam.py with PYAUTO_TEST_MODE=2, bypassing the sampler before _fit — a vacuous pass. After it, that run really executes the JAX path, so a backend regression is caught within a week on the weekly validation, instead of only at release time. That still satisfies the intent ("it should not stay bypassed if running it is what ensures the bug is caught") — it just is not per-PR.

It also means the searches/BlackJAXNUTS precedent I copied is itself about the full-script-set runs, not the per-PR gate. My reading of it was right; my description of what it gates was not.

Nothing about the diff changes. All four overrides remain correct and verified, and the fix for the blocking failure (the release profile) is unaffected — that was always the load-bearing part.

The open judgement I am NOT deciding unilaterally: getting genuine per-PR coverage would mean adding searches/MultiStartAdam.py to smoke_tests.txt. That is promoting an integration script into the deliberately-small curated subset, which cuts against the standing "smoke tests are a small curated subset — don't mass-promote" preference. Worth a separate decision on its own merits; say the word and I will file it.

CI status on this PR: smoke (3.12) and smoke (3.13) both SUCCESS — though per the above, those runs did not exercise MultiStartAdam and are not evidence for this change.

[Jammy2211]

Shipped — PR #45 squash-merged as ad95dd8 on main.

Landed: three set: { PYAUTO_DISABLE_JAX: "0" } overrides in env_vars_release.yaml (MultiStartAdam, Dynesty_jax, Nautilus_jax) + one unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX] in env_vars.yaml. Config only; no source, no API.

Verified before merge: all three scripts exit 0 with JAX enabled (MultiStartAdam recovers 50.156/25.196/9.858 vs truth 50/25/10); env resolution driven through autobuild.env_config.build_env_for_script over the real configs for every scripts/searches/*.py, confirming no substring over-match.

The real confirmation is still pending: tonight's nightly-release reaching Stage 3 clean for the first time in six nights. Closing on the merge, but that run is the actual proof — will reopen if the shard is still red.

Carried forward: the smoke-profile change buys a weekly net (Heart's workspace-validation.yml mode=smoke), not a per-PR one — see the correction comment above. Genuine per-PR coverage would require promoting MultiStartAdam.py into the curated smoke_tests.txt, which is a deliberate open decision, not an oversight.