docs(jax): correct postprocessing_vectorize default in sample_jax_nuts docstring

[xodn348]

Description

Fixes a docstring inaccuracy in sample_numpyro_nuts / sample_blackjax_nuts (the shared parameter section in pymc/sampling/jax.py).

The parameter description for postprocessing_vectorize states default "scan", but the actual runtime default is "vmap". When postprocessing_vectorize=None (i.e., the caller uses the default), the function body sets it to "vmap":

else:
    postprocessing_vectorize = "vmap"

This mismatch led a user (#8238) to remove their explicit postprocessing_vectorize="scan" argument (after reading that "scan" was the default) and accidentally switch to "vmap", causing an out-of-memory error.

The one-character change corrects default "scan" → default "vmap" so the docstring matches the actual behaviour.

Related Issue

• Closes postprocessing_vectorize default does not match docs/warnings #8238

Checklist

• Checked that the pre-commit linting/style checks pass (pre-commit run --files pymc/sampling/jax.py — all hooks passed)
• Included tests that prove the fix is effective or that the new feature works — N/A: this is a docstring-only correction with no behaviour change; the existing code path is fully tested by the existing JAX test suite
• Added necessary documentation (docstrings and/or example notebooks) — the fix IS the docstring correction
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• Documentation

Local verification

$ pre-commit run --files pymc/sampling/jax.py
check for merge conflicts............................................................Passed
debug statements (python)............................................................Passed
fix end of files.....................................................................Passed
don't commit to branch...............................................................Passed
trim trailing whitespace.............................................................Passed
check blanket noqa...................................................................Passed
check blanket type ignore............................................................Passed
check for not-real mock methods......................................................Passed
use logger.warning(..................................................................Passed
type annotations not comments........................................................Passed
no unicode replacement chars.........................................................Passed
Apply Apache 2.0 License.............................................................Passed
ruff check...........................................................................Passed
ruff format..........................................................................Passed

$ ruff check pymc/sampling/jax.py
All checks passed!

$ ruff format --check pymc/sampling/jax.py
1 file already formatted

=== LOCAL_TEST_PASSED ===

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #32600982 | 📁 Comparing 62d8d38 against latest (151672a)

  🔍 Preview build  

3 files changed

± glossary.html
± api/generated/pymc.sampling.jax.sample_blackjax_nuts.html
± api/generated/pymc.sampling.jax.sample_numpyro_nuts.html