Raise on unrecognized sample() kwargs for external NUTS samplers

[sangaraju1988]

Description

#8352 asked for a way to turn off init jitter when using an external NUTS sampler. Half of that got fixed already, as a side effect of #8369 — nuts={"jitter": False} works fine on main right now.

The part still broken is the calling style from before that refactor. pm.sample(nuts_sampler="numpyro", jitter=False) — jitter passed at the top level instead of nested in nuts={} — does nothing. No error, no warning. It just samples with jitter on anyway.

Reason: _sample_external_nuts() takes **kwargs as a catch-all, and nothing in the function ever reads it. Whatever lands there just evaporates. So this fixes the catch-all itself, not just the jitter case — it now raises a TypeError naming whatever unexpected keys showed up, telling you to use nuts={...} instead.

Added two tests. test_external_nuts_sampler_rejects_stray_top_level_kwargs is the one that actually covers this change — confirmed it fails with DID NOT RAISE on unpatched main and passes with the fix. test_jax_sampler_jitter_kwarg_routing covers the nuts={"jitter": False} path, which already worked before this PR (that's #8369's doing) — including it since it wasn't tested anywhere before now.

Ran test_jax.py, test_mcmc_external.py, and test_mcmc.py locally. Full pass on the first two. test_mcmc_external.py has some pre-existing flakiness unrelated to this — a couple of blackjax tests occasionally hit a pmap/jax version issue, and test_step_args has a statistical assertion that isn't always stable. Both reproduce identically on unpatched main, so not from this change.

Related Issue

• Closes BUG: pm.sample can't disable init-value jitter for the external JAX samplers (numpyro/blackjax) #8352

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)

Type of change

• Bug fix

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #34213935 | 📁 Comparing e46b453 against latest (da8fc47)

  🔍 Preview build  

2 files changed

± glossary.html
± _modules/pymc/sampling/mcmc.html