From 615b22920af233524343ab858a5549dea2ac1081 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 15 Jul 2026 08:45:20 +0100 Subject: [PATCH 1/2] fix(config): enable JAX for the JAX-native searches scripts (#44) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The mode=release profile (env_vars_release.yaml) pins PYAUTO_DISABLE_JAX=1 alongside PYAUTO_TEST_MODE=0, so searches scripts really run but on the NumPy backend. MultiStartAdam is JAX-native and hard-raises when the analysis is not JAX-traceable, so it has failed the autofit_test/searches shard — and with it nightly-release Stage 3 — every night since it landed in #43. Add the missing `set: { PYAUTO_DISABLE_JAX: "0" }` overrides, following the searches/BlackJAXNUTS precedent. Dynesty_jax and Nautilus_jax get the same treatment: they declare use_jax=True but were also absent from the overrides, so they were validating the NumPy path while presenting as JAX tests — passing vacuously, since unlike MultiStartAdam they carry no backend guard. Also stop the smoke profile bypassing MultiStartAdam. Its PYAUTO_TEST_MODE=2 skips the sampler before _fit runs, so the script passed without executing its JAX path at all — which is why the per-PR gate could not catch this and it surfaced only in mode=release. Running it costs ~6s. Verified with JAX enabled: MultiStartAdam exit 0 (recovers 50.156/25.196/9.858 vs truth 50/25/10), Dynesty_jax exit 0, Nautilus_jax exit 0, and MultiStartAdam under the smoke-resolved env exit 0. Env resolution driven through autobuild.env_config.build_env_for_script over the real configs for every scripts/searches/*.py: Nautilus.py and DynestyStatic/Dynamic.py correctly stay at DISABLE_JAX=1, confirming the substring patterns do not over-match. Co-Authored-By: Claude Opus 4.8 --- config/build/env_vars.yaml | 9 +++++++++ config/build/env_vars_release.yaml | 14 ++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/config/build/env_vars.yaml b/config/build/env_vars.yaml index 0b54277..d2e13c6 100644 --- a/config/build/env_vars.yaml +++ b/config/build/env_vars.yaml @@ -56,3 +56,12 @@ overrides: # at the script's already-cheap defaults (num_warmup=200, num_samples=500). - pattern: "searches/BlackJAXNUTS" unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX] + # MultiStartAdam is a JAX-native gradient MAP search that hard-raises when the + # analysis is not on the JAX backend. Under the smoke defaults it never + # reaches that guard: PYAUTO_TEST_MODE=2 bypasses the sampler before + # `_fit` runs, so the script passes without executing its JAX path at all — + # which is exactly why the backend regression it was written to catch got + # through this gate and only failed in mode=release. Run it for real (a + # jitted 3-parameter fit, a few seconds) so the per-PR gate can see it. + - pattern: "searches/MultiStartAdam" + unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX] diff --git a/config/build/env_vars_release.yaml b/config/build/env_vars_release.yaml index 0609d68..aa630f6 100644 --- a/config/build/env_vars_release.yaml +++ b/config/build/env_vars_release.yaml @@ -52,3 +52,17 @@ overrides: # already satisfied by defaults. - pattern: "searches/BlackJAXNUTS" set: { PYAUTO_DISABLE_JAX: "0" } + # MultiStartAdam is a JAX-native gradient MAP search (jax.vmap of + # jax.value_and_grad over the autofit Fitness). Unlike the samplers below it + # hard-raises when the analysis is not on the JAX backend, so the disabling + # default turned it into a release-validation failure the moment it landed. + - pattern: "searches/MultiStartAdam" + set: { PYAUTO_DISABLE_JAX: "0" } + # The *_jax sampler scripts exist to validate the JAX likelihood path. They + # carry no hard backend guard, so under the disabling default they silently + # validated the NumPy path and passed vacuously — the same failure mode the + # jax_assertions/ note above describes, minus the loud failure. + - pattern: "searches/Dynesty_jax" + set: { PYAUTO_DISABLE_JAX: "0" } + - pattern: "searches/Nautilus_jax" + set: { PYAUTO_DISABLE_JAX: "0" } From cb1b8f1837136247b98414b42f1da9fb4d20c15a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 15 Jul 2026 08:49:55 +0100 Subject: [PATCH 2/2] docs(config): correct what the smoke profile actually gates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The comment added in 615b229 claimed running MultiStartAdam under this profile lets "the per-PR gate" see it. That is wrong, and the wrong claim would have landed in main. Two different things are called smoke here: 1. The per-PR gate is .github/workflows/smoke_tests.yml, which runs the curated smoke_tests.txt subset. Neither MultiStartAdam nor BlackJAXNUTS is in that list, so this profile never touches it. 2. env_vars.yaml is the `smoke` env profile consumed by PyAutoHeart's workspace-validation.yml at mode=smoke, which runs every script weekly. The override lands on (2). The change is still worth keeping — the weekly sweep now really executes the JAX path instead of bypassing it, so a backend regression is caught within a week rather than at release — but on a weekly cadence, not per-PR. Correct the comment to say so, and record that the BlackJAXNUTS precedent above is likewise about the full-script-set runs. No behaviour change: comment only. Co-Authored-By: Claude Opus 4.8 --- config/build/env_vars.yaml | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/config/build/env_vars.yaml b/config/build/env_vars.yaml index d2e13c6..6034aad 100644 --- a/config/build/env_vars.yaml +++ b/config/build/env_vars.yaml @@ -57,11 +57,14 @@ overrides: - pattern: "searches/BlackJAXNUTS" unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX] # MultiStartAdam is a JAX-native gradient MAP search that hard-raises when the - # analysis is not on the JAX backend. Under the smoke defaults it never - # reaches that guard: PYAUTO_TEST_MODE=2 bypasses the sampler before - # `_fit` runs, so the script passes without executing its JAX path at all — - # which is exactly why the backend regression it was written to catch got - # through this gate and only failed in mode=release. Run it for real (a - # jitted 3-parameter fit, a few seconds) so the per-PR gate can see it. + # analysis is not on the JAX backend. Under this profile's defaults it never + # reaches that guard: PYAUTO_TEST_MODE=2 bypasses the sampler before `_fit` + # runs, so the script passes without executing its JAX path at all — the + # backend regression it exists to catch went unnoticed here and surfaced only + # under mode=release. Run it for real (a jitted 3-parameter fit, ~6s) so the + # weekly mode=smoke workspace-validation sweep, which runs every script under + # this profile, actually exercises it. Note this profile does NOT feed the + # per-PR gate — that is the curated smoke_tests.txt list, which contains + # neither this script nor BlackJAXNUTS. - pattern: "searches/MultiStartAdam" unset: [PYAUTO_TEST_MODE, PYAUTO_DISABLE_JAX]