Skip to content

fix: populate NUTS samples_info keys under test-mode bypass#1260

Merged
Jammy2211 merged 1 commit intomainfrom
feature/fix-smoke-jax-and-mcmc
May 8, 2026
Merged

fix: populate NUTS samples_info keys under test-mode bypass#1260
Jammy2211 merged 1 commit intomainfrom
feature/fix-smoke-jax-and-mcmc

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

  • Adds _test_mode_samples_info() hook on AbstractSearch so subclasses can declare which diagnostic keys their samples_info_from() would normally populate. Hook return is merged into the bypass samples_info so downstream code reading per-sampler diagnostics doesn't KeyError.
  • Overrides the hook in BlackJAXNUTS to return num_warmup, num_samples, num_chains, ess_min, ess_per_param, mean_acceptance, n_divergent, n_logl_evals, total_walkers, total_steps with NaN/0 placeholders.
  • Unblocks autofit_workspace/scripts/searches/mcmc.py smoke (NUTS section) under PYAUTO_TEST_MODE=2.

Test plan

  • Local: PYAUTO_TEST_MODE=2 python autofit_workspace/scripts/searches/mcmc.py runs to completion, prints ESS (min over dims): nan / 0 etc.
  • CI: Tests workflow on this PR
  • CI: autofit_workspace smoke on main after merge

🤖 Generated with Claude Code

`PYAUTO_TEST_MODE=2` skips the sampler in `_fit_bypass_test_mode`,
producing a samples_info stub with only `total_iterations`, `time`,
and `log_evidence`. Tutorial and downstream code that reads NUTS
diagnostics (`ess_min`, `num_samples`, `mean_acceptance`,
`n_divergent`, `n_logl_evals`) crashed with `KeyError`.

Add a `_test_mode_samples_info()` hook on `AbstractSearch`, merged
into the bypass `samples_info`. Override in `BlackJAXNUTS` to return
the diagnostic keys with NaN/0 placeholders — bypass mode never
ran the sampler, so honest empties propagate to the tutorial print.

Fixes the `autofit_workspace/scripts/searches/mcmc.py` smoke test.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
@Jammy2211 Jammy2211 merged commit 5d175eb into main May 8, 2026
3 checks passed
@Jammy2211 Jammy2211 deleted the feature/fix-smoke-jax-and-mcmc branch May 8, 2026 15:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant