test(searches): validate MultiStartAdam recovers truth basin under JAX#43
Merged
Merged
Conversation
Adds scripts/searches/MultiStartAdam.py — the JAX cross-backend end-to-end validation for af.MultiStartAdam (PyAutoFit#1369), asserting truth-basin recovery on the 1D Gaussian dataset (centre/normalization/sigma = 50/25/10) via the standard search.fit() API with a use_jax=True analysis. Kept out of PyAutoFit's NumPy-only unit suite by design. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Jammy2211
added a commit
that referenced
this pull request
Jul 15, 2026
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. Note what this does and does not gate: env_vars.yaml feeds Heart's workspace-validation.yml at mode=smoke, which runs every script weekly — NOT the per-PR gate, which is the curated smoke_tests.txt list containing neither this script nor BlackJAXNUTS. So this buys a weekly net on the JAX path, not a per-PR one. 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. Closes #44 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The JAX cross-backend end-to-end validation for the new
af.MultiStartAdamsearch (PyAutoFit#1369, PR#1370 merged). Addsscripts/searches/MultiStartAdam.py, which runs the multi-start gradient MAP search on the 1D Gaussian dataset via the standardsearch.fit()API with ause_jax=Trueanalysis, and asserts recovery of the known truth basin (centre/normalization/sigma = 50/25/10).This is the remaining Phase-1 leg of the promotion. It lives here rather than in PyAutoFit's unit suite because library unit tests stay NumPy-only; the real JAX/optax path is exercised here.
Test Plan
python scripts/searches/MultiStartAdam.py→ recoveredcentre=50.156, normalization=25.196, sigma=9.858; assertion passes (jax 0.10.2 + optax 0.2.5)Notes
Nautilus_jax.pypattern (enable_pytrees()+register_model(),af.ex.Analysis(use_jax=True),NullPaths).smoke_tests.txt— the*_jaxsearch scripts are not part of the curated smoke subset.Generated by the PyAutoLabs agent workflow.