add Dynesty_jax.py — JAX-jitted likelihood example for DynestyStatic#19
Merged
add Dynesty_jax.py — JAX-jitted likelihood example for DynestyStatic#19
Conversation
Companion to Nautilus_jax.py. Exercises the new jit-only JAX path on DynestyStatic introduced in PyAutoFit#1243. Same 1D Gaussian dataset, NullPaths search, ~3.5s runtime. Verifies the path executes end-to-end and emits the expected `JAX: Applying jit to likelihood function` log line — not the same as Nautilus's vmap+jit path because dynesty 2.1.5 has no `vectorized` parameter and calls the likelihood one sample at a time. Co-Authored-By: Claude Opus 4.7 (1M context) <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
scripts/searches/Dynesty_jax.py— companion toNautilus_jax.py, exercises the JAX-jit path onaf.DynestyStatic.use_jax_jitplumbing in PyAutoFit#1243 (just merged) executes end-to-end on the 1D Gaussian dataset.Why
Companion test for PyAutoFit#1243, which added
jax.jitlikelihood support to dynesty. Without a runnable example inautofit_workspace_test, there's no integration check on the new JAX-jit path beyond the library unit tests (which deliberately don't import jax).dynesty differs from Nautilus on this path: dynesty 2.1.5's
NestedSamplerhas novectorizedparameter, so vmap-batching doesn't apply. The script usesjax.jitonly — JAX's compiled-function cache reuses the compiled likelihood across dynesty's one-sample-at-a-time call pattern.Verification (local)
JAX: Applying jit to likelihood functionlog line emitted (the analogue of Nautilus'sJAX: Applying vmap and jit...line)Running Dynesty with JAX-jitted likelihood (single CPU, no pool).from PyAutoFit#1243's new log branchlog_Z ≈ -54.0, dlogz ≈ 0.014(tolerance 0.5)nlive=30, dlogz=0.5Related
Nautilus_jax.py, the precursorTest plan
python3 scripts/searches/Dynesty_jax.pyruns end-to-end without errorJAX: Applying jit to likelihood functionRunning Dynesty with JAX-jitted likelihoodNautilus_jax.pystill passes (vmap regression check, log_Z ≈ -54.13)🤖 Generated with Claude Code