Every way of sampling a model is now a Sampler with two equivalent
surfaces: a configured object for pm.sample(sampler=...) and a flat
entry point (pm.nutpie.nuts.sample(...)), under the rule "constructor
takes algorithm configuration, pm.sample takes run configuration".
This is the refactor-only unit agreed in pymc#8353; new blackjax
algorithm support builds on it in a follow-up.
- Sampler ABC: constructors are algorithm configuration only (nothing
model-bound, nothing compiled); the model and the shared run
configuration arrive at sample time via sample_from_init, whose
signature is kept in sync with pm.sample by a test. The method is
deliberately not named `sample` to distinguish it from the flat
entry points. ExternalSampler adds only a dependency check.
- StepSampler wraps the native machinery (step assignment,
CompoundStep, the python loop) as a regular sampler. The engine was
extracted verbatim from pm.sample into _sample_with_step_methods and
is shared by both, so pm.sample() and
pm.sample(sampler=pm.StepSampler()) produce identical draws
(asserted by a test). NUTS initialization configuration (init,
n_init, jitter_max_retries) lives on its constructor.
- Nutpie, NumpyroNUTS and BlackjaxNUTS wrap the existing external
NUTS implementations, each documenting per run argument whether it
is honored, reinterpreted, warned about, or rejected (e.g. a
non-jax compile mode raises for the jax samplers while nutpie
consumes it). Exposed as pm.nutpie.nuts, pm.numpyro.nuts and
pm.blackjax.nuts.
- pm.sample grows sampler=, forwarding only run configuration and
rejecting NUTS-specific arguments; nuts_sampler= and explicit
init/n_init/jitter_max_retries now emit DeprecationWarning (not yet
FutureWarning) pointing at their new homes.
- Pin pyarrow<25 in the alternative_backends CI job (transitive
nutpie dependency that segfaults XLA CPU compilation in-process).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
What
The refactor-only unit requested in #8353 review: every existing way of sampling a model becomes a
Samplerwith two equivalent surfaces, no new sampling capability is added. The generic blackjax algorithm support in #8353 will be rebased on top of this once merged.How
Following the class sketch from the #7880/#8353 discussion:
SamplerABC (pymc/sampling/samplers/base.py): constructors are algorithm configuration only — nothing model-bound, nothing compiled. The model and the shared run configuration arrive at sample time viasample_from_init, whose signature is exactlypm.sample's run subset (kept in sync by a test, no**kwargsescape hatch). Deliberately not namedsampleso it can't be confused with the flat entry points.ExternalSampleris implementation inheritance only: apackage/version_rangedependency check with an install hint.SamplerEntrymakes the two surfaces one implementation:pm.nutpie.nuts(...)returns a configured sampler for the funnel;pm.nutpie.nuts.sample(...)is the flat call, defined in terms of the first so they cannot disagree.StepSampler: the native machinery (step assignment,CompoundStep, the python loop) as a regular sampler. The engine was extracted verbatim frompm.sampleinto_sample_with_step_methodsand is shared by both paths — a test assertspm.sample()andpm.sample(sampler=pm.StepSampler())produce identical draws. NUTS initialization configuration (init,n_init,jitter_max_retries) lives on its constructor, which is also the first step toward fixing init_nuts not used when CompoundStep is used #5658.Nutpie,NumpyroNUTS,BlackjaxNUTSwrap the existing external NUTS implementations. Each documents an explicit disposition per run argument: nutpie consumescompile_kwargs/backend(numba/jax) and honorsdiscard_tuned_samples; the jax pair raises on a non-jax compile mode instead of silently ignoring it, and warns ondiscard_tuned_samples=False/keep_warning_stat.pm.sample(sampler=...)forwards only the run configuration and rejects the NUTS-specific arguments with pointers to their new homes.Deprecations (DeprecationWarning, not FutureWarning yet)
nuts_sampler=→pm.sample(sampler=pm.nutpie.nuts())/pm.numpyro.nuts()/pm.blackjax.nuts()/pm.StepSampler(), or the flat entry points.init/n_init/jitter_max_retries→pm.sample(sampler=pm.StepSampler(init=...)).pm.sample()is unaffected.The acceptance criteria from the discussion hold by construction: the internal NUTS is expressible as a regular sampler through the same interface with no special-casing in
pm.sample, and everynuts_sampler=string value has a class equivalent, so the argument is removable at the end of the deprecation cycle.Also included
pyarrow<25pin in thealternative_backendsCI job — pyarrow 25.0.0 (unpinned transitive dependency of thenutpie@maininstall) segfaults XLA CPU compilation when loaded in the same process; without this pin the job fails on every PR (root-cause analysis in #8353).Tests
tests/sampling/test_samplers.py: StepSampler via funnel/flat/explicit-step plus the identical-draws equivalence test, all three external NUTS samplers via both surfaces, funnel forwarding/rejections, the signature-sync test across all sampler classes, deprecation warnings (and their absence for defaults), and dependency-check errors. Fulltests/sampling/test_mcmc.pypasses unchanged (106 tests).🤖 Generated with Claude Code