fix(config): enable JAX for the JAX-native searches scripts

[Jammy2211]

Fixes #44.

What broke

nightly-release has been red five nights running (2026-07-11 → 2026-07-15). Stage 3 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.

Every other shard in that run is green. This is not the shifting perf-flake tail from the release-tail epic — it is deterministic, and it is fallout from the multi-start gradient search promotion (PyAutoFit #1370, this repo's #43).

Root cause

MultiStartAdam._fit hard-guards on the backend:

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 gate — pins PYAUTO_TEST_MODE: "0" (so _fit really runs, no bypass) and PYAUTO_DISABLE_JAX: "1" (so use_jax=True is silently flipped to False). The guard fires. The profile re-enables JAX for exactly two patterns — jax_assertions/ and searches/BlackJAXNUTS — and searches/MultiStartAdam was never added.

autofit_workspace/scripts/searches/mle.py uses the same search and passes only because that workspace's release profile pins PYAUTO_DISABLE_JAX: "0". The autolens/autogalaxy guides construct the search but never .fit(), so they never reach the guard.

Changes

Config only — no script logic, no library source, no public API.

config/build/env_vars_release.yaml — three set: { PYAUTO_DISABLE_JAX: "0" } overrides, following the searches/BlackJAXNUTS precedent (this profile's header documents that overrides only ever set:, never unset:, a var defaults pins):

• searches/MultiStartAdam — the blocking fix.
• searches/Dynesty_jax, searches/Nautilus_jax — these declare use_jax=True but were also missing from the overrides, so they validated the NumPy path while presenting as JAX tests. Unlike MultiStartAdam they carry no backend guard, so they passed vacuously and nothing failed loudly — the same failure mode the profile's own jax_assertions/ note describes.

config/build/env_vars.yaml (the smoke env profile) — one unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX] for searches/MultiStartAdam. Its PYAUTO_TEST_MODE=2 bypasses the sampler before _fit, so the script passed without ever executing its JAX path. Running it costs ~6s. Scoped to MultiStartAdam; the nested samplers stay bypassed under this profile on cost grounds.

What this profile does and does not gate (an earlier revision of this PR body got this wrong — see the correction comment). env_vars.yaml feeds PyAutoHeart's workspace-validation.yml at mode=smoke, which runs every script weekly (Mondays 03:00 UTC). It does not feed the per-PR gate — that is .github/workflows/smoke_tests.yml running the curated smoke_tests.txt list, which contains neither MultiStartAdam nor BlackJAXNUTS. So this override buys a weekly net on the JAX path, not a per-PR one. Genuine per-PR coverage would mean promoting the script into smoke_tests.txt — a separate decision, against the standing "keep the curated subset small" preference, deliberately not taken here.

Verification

All run on this branch 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 (smoke-resolved env, both vars unset)	exit 0 (~6s)

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

Resolver check. _pattern_matches substring-matches against the full path including extension, so an under-specified pattern silently over-matches (searches/Nautilus would also catch Nautilus_jax.py). Drove autobuild.env_config.build_env_for_script over the real configs for every scripts/searches/*.py:

===== 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 — no over-match.

Ship-gate legs

1. Tests — n/a: no library test dir in this repo; the scripts are the tests, and every affected one was executed directly (above).
2. Smoke — the changed surface is the searches shard; env resolution verified against the real configs and the three JAX scripts run end-to-end. The PR's own smoke (3.12/3.13) checks pass but do not exercise MultiStartAdam (it is not in smoke_tests.txt), so they are not evidence for this change; the direct runs above are.
3. Review — 23 lines of config + rationale comments.
4. Heart — YELLOW score 55 at ship time. Its first reason, "workspace validation not passing", is what this PR fixes.

Human signed off on the ship in-session (autonomy: supervised, the bug cap). Merge stays human.

🤖 Generated with Claude Code

[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.