Skip to content

add Nautilus_jax.py — JAX-jitted Nautilus integration test#17

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/nautilus-jax-example
Apr 30, 2026
Merged

add Nautilus_jax.py — JAX-jitted Nautilus integration test#17
Jammy2211 merged 1 commit into
mainfrom
feature/nautilus-jax-example

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Adds scripts/searches/Nautilus_jax.py, a JAX-jitted variant of the existing Nautilus.py integration test. Exercises the end-to-end JAX path: enable_pytrees() + register_model(model) + af.ex.Analysis(use_jax=True) + af.Nautilus(...) against the 1D Gaussian dataset. Confirms Fitness._vmap (which wraps the likelihood in jax.vmap(jax.jit(...))) and af.Nautilus's auto-routing to fit_x1_cpu with vectorized=True work end-to-end against af.Model(af.ex.Gaussian).

Local run: search finishes (log_Z ≈ -64.15, n_eff = 366) in ~24s. The JAX: Applying vmap and jit to likelihood function log line in stdout confirms the JIT path is exercised, not silently bypassed.

Scripts Changed

  • scripts/searches/Nautilus_jax.py (new) — JAX-jitted Nautilus integration test, NullPaths (n_live=50, n_eff=100).

Upstream PR

PyAutoLabs/PyAutoFit#1240 (merged) — fixes the bare-Model xp fallback in af.ex.Analysis.model_data_1d_from that this script depends on.

Test Plan

  • python scripts/searches/Nautilus_jax.py runs to completion under $WT_ROOT/activate.sh
  • No failed/scripts__searches__Nautilus_jax.py.txt log produced by bash run_all_scripts.sh

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Workspace fixes are now up: PyAutoLabs/autofit_workspace#47 (the actual breaking-case fix for the inline tutorial classes) and Jammy2211/autofit_workspace_developer#9 (defensive parity for 7 dev scripts). All three can land together once #47's smoke CI is green.

@Jammy2211 Jammy2211 merged commit aa5aa17 into main Apr 30, 2026
4 checks passed
@Jammy2211 Jammy2211 deleted the feature/nautilus-jax-example branch April 30, 2026 14:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant