test(searches): JAX validation for af.MultiStartProdigy#58
Merged
Conversation
Add scripts/searches/MultiStartProdigy.py, mirroring MultiStartAdam.py but running the learning-rate-free Prodigy rule (no learning_rate). Asserts truth-basin recovery on the 1D Gaussian dataset and exercises the parts the NumPy-only unit suite cannot: optax.contrib resolution, per-start vmapped optimizer state, and the apply_if_finite guard on the real JAX backend. Upstream: PyAutoLabs/PyAutoFit#1398. 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
Adds the JAX end-to-end validation script for
af.MultiStartProdigy, the learning-rate-free multi-start gradient MAP search promoted in the upstream library PR (Phase 1 of the multi-start gradient v2 promotion, autolens_workspace_developer#101).Mirrors the existing
MultiStartAdam.pybut runs Prodigy with nolearning_rate, and asserts it recovers the 1D Gaussian truth basin (centre, normalization, sigma = 50, 25, 10). This is the JAX cross-backend check that deliberately lives here rather than in PyAutoFit's NumPy-only unit suite: it exercises the parts the unit suite cannot —optax.contribrule resolution, the per-startjax.vmaped optimizer state (Prodigy's globaldestimate kept independent across starts), and theoptax.apply_if_finiteguard — all on the real JAX backend.Upstream PR
PyAutoLabs/PyAutoFit#1398 (
pending-release) — must merge before this workspace PR (library-first merge gate).Scripts Changed
Added
scripts/searches/MultiStartProdigy.py— end-to-end JAX validation ofaf.MultiStartProdigy; asserts learning-rate-free truth-basin recovery on the 1D Gaussian dataset.No existing scripts change (the library API change is additive and backward-compatible; existing
MultiStart*scripts are unaffected). Not added tosmoke_tests.txt— it follows theMultiStartAdam.pyprecedent as a full-sweep validation script, keeping the curated smoke subset small.Test Plan
python scripts/searches/MultiStartProdigy.py— passes:Recovered: centre=50.156, normalization=25.196, sigma=9.858, no learning rate.Generated by the PyAutoLabs agent workflow.