feat: re-mainline af.NSS as a first-class search on mainline blackjax ≥1.6 - #1492
Merged
Conversation
… >=1.6 Restores the NSS nested slice sampler removed in #1356/#1357 (the removal reason -- unshippable git-fork pins -- is gone: nested sampling merged into mainline blackjax in release 1.6). Implementation restored verbatim from the autofit_workspace_developer/searches/nss/ stash with the mainline port: - from nss.ns import log_weights, finalise -> blackjax.ns.utils; the import guard now checks blackjax >= 1.6 instead of the fork packages. - SliceInfo eval counter: fork num_steps -> mainline num_expansions. - Chunked GPU-memory path recomposed from mainline public helpers (build_slice_kernel / slice_constrained_step / build_adaptive_kernel / live_covariance) since mainline build_kernel has no update_strategy kwarg; _chunked_update.py unchanged (already byte-equivalent to mainline's update_with_mcmc_take_last bar the vmap->lax.map swap). - Tests restored verbatim (one ImportError match-string update); NSS re-exported from autofit/__init__.py; optional extra blackjax floor bumped 1.2.0 -> 1.6.2 (no cap existed). Validation (CPU): 17/17 nss tests pass from the checkout; end-to-end 2D analytic toy through search.fit gives logZ -4.558 +/- 0.078 vs analytic -4.605 with truth recovered, and the chunked path is bit-identical to unchunked at fixed seed (same logZ/evals/ESS). Closes #1491. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018dV9h5h76hZ6trGZKWJFHe
af.NSS samples through inline JAX closures, never through Fitness -- its Fitness exists only for AbstractNest.perform_update post-fit work, so forwarding iterations_per_quick_update there would wire the cadence to a dead path. Quick updates fire through the search's own outer loop (_fire_quick_update), so the announced cadence is real. The wiring test postdates the search's removal, which is why the restored stash tripped it. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018dV9h5h76hZ6trGZKWJFHe
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
2 participants
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.
Closes #1491. Restores the
af.NSSnested slice sampler removed in #1356/#1357. The removal reason — unshippable git-fork pins (handley-lab/blackjax+yallup/nss) — is gone: nested sampling merged into mainline blackjax 1.6 (current 1.6.2), so this is now a plain-dependency search, exactly the condition the stash README's re-mainlining checklist was waiting for.What this restores
autofit/non_linear/search/nest/nss/(search + samples + the chunked GPU-memory builders) andtest_autofit/non_linear/search/nest/nss/, verbatim from theautofit_workspace_developer/searches/nss/stash, with the mainline port:from nss.ns import log_weights, finalise→blackjax.ns.utils; the import guard now checksblackjax >= 1.6instead of the fork packages.SliceInfo.num_steps→ mainlinenum_expansions(+num_shrink, unchanged).build_kernelhas noupdate_strategy=kwarg (the fork's seam), so_chunked_nss.pyrecomposes the kernel from mainline public helpers (build_slice_kernel/slice_constrained_step/build_adaptive_kernel/live_covariance)._chunked_update.pyneeded no logic change — its body is byte-equivalent to mainline'supdate_with_mcmc_take_lastbar the final vmap→lax.mapswap.NSSre-exported fromautofit/__init__.py;pyproject.tomloptional-extra blackjax floor bumped>=1.2.0→>=1.6.2(no cap existed).Validation (CPU, from the checkout)
test_autofit/non_linear/search/nest/nss/).search.fit(2D Gaussian × uniform box, analytic logZ = −4.605): logZ = −4.558 ± 0.078, posterior median (1.008, −0.990) vs truth (1, −1).chunk_size=7produces bit-identical logZ / evals / ESS to the unchunked path at fixed seed — the strongest parity check for the recomposed builder.Context
Requested by the human 2026-08-18 as part of the inference programme (autolens_profiling
results/notes/inference/PROGRAMME.md); the programme's Gate A still decides whether NSS becomes the GPU nested-sampling baseline — this PR only restores the search. Follow-ups tracked separately: the workspace tutorial section (stash checklist item 5) and the Phase 2 NSS re-benchmarking campaign.🤖 Generated with Claude Code
https://claude.ai/code/session_018dV9h5h76hZ6trGZKWJFHe
Generated by Claude Code