Generic Blackjax external sampler: run any blackjax MCMC algorithm via pm.sample#8353
Generic Blackjax external sampler: run any blackjax MCMC algorithm via pm.sample#8353twiecki wants to merge 18 commits into
Conversation
Implements pm.sample(external_sampler=pm.external.Blackjax(...)), a generic driver that can run any blackjax MCMC algorithm on a PyMC model, building on the ExternalSampler API drafted in #7880. - Algorithms are resolved dynamically against the blackjax namespace (validated via GenerateSamplingAPI), so new blackjax algorithms work without changes on the PyMC side. - Keyword arguments are routed to the algorithm factory or the adaptation procedure by signature introspection, with a discoverable get_kwargs(); unknown/missing parameters raise informative errors at construction time. - A small curated table maps algorithms to their default adaptation: window adaptation for the HMC family, the bespoke mclmc_find_L_and_step_size tuner for MCLMC, and plain burn-in for static-parameter algorithms (mala, barker, rmh, ...). - Sampler stats are flattened generically from the blackjax Info namedtuples and renamed to ArviZ conventions. - The idata assembly of sample_jax_nuts is extracted into a reusable _postprocess_and_build_idata helper shared with the new driver. SMC, SGMCMC, cross-chain adaptation (MEADS/ChEES) and variational algorithms are out of scope and raise errors pointing at alternatives. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Documentation build overview
29 files changed ·
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8353 +/- ##
========================================
Coverage 91.73% 91.74%
========================================
Files 128 131 +3
Lines 20672 20970 +298
========================================
+ Hits 18963 19238 +275
- Misses 1709 1732 +23
🚀 New features to boost your workflow:
|
Importing blackjax at test collection time initialized the jax backend before pymc.sampling.jax could set XLA_FLAGS, leaving a single host device and crashing every chain_method="parallel" test in the session. Import pymc.sampling.jax first in the test module, and make the Blackjax sampler fall back to vectorized chains with a warning when there are fewer devices than chains. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Module-type names cannot be re-exported through pymc.sampling.__all__ (test_root_namespace excludes modules); import it at the root like pm.math and pm.gp instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
please address #7880 (comment) |
Address review feedback on the external sampler API: - pm.external.blackjax.<algorithm>() factories (e.g. pm.external.blackjax.mclmc()) via module __getattr__/__dir__, enumerated dynamically from the installed blackjax version, as a discoverable per-algorithm alternative to the algorithm string. - Every pm.sample argument now has an explicit disposition on the external sampler path: forwarded, mapped (init -> jitter/no-jitter, fixing the semantics discussed in #8352), warned-and-ignored (non-mappable init strings, discard_tuned_samples=False, keep_warning_stat, backend/compile_kwargs), or rejected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Both points from that comment are now addressed in cb0bac8: 1. Per-algorithm entries instead of algorithm strings. pm.sample(external_sampler=pm.external.blackjax.mclmc())
pm.sample(external_sampler=pm.external.blackjax.nuts(target_accept=0.9))
pm.sample(external_sampler=pm.external.blackjax.mala(step_size=0.01))Implemented via module 2. The
The table is documented in the I agree this doesn't dissolve your underlying discomfort with the 🤖 Generated with Claude Code |
|
For reviewers wondering about the red CI runs and revert/reapply commits on Jul 10: the Symptom: deterministic native segfault inside XLA CPU compilation ( Root cause: Evidence:
In this PR: the job now installs 🤖 Generated with Claude Code |
Disagree quite a lot with the conclusion. jitter and init are very specific to nuts, backend shouldn't be ignored (e.g., nutpie takes it). From this iteration I have 0 confidence on the bot wisdom, so next iteration try to make an argument for each. Not all external samplers are NUTS and not all support a single backend. |
This reverts commit 00710a6.
Restores pm.sample(external_sampler=...) (reverted in the previous commit) but relocates all argument interpretation out of the shared path, per review: the meaning of pm.sample's arguments is sampler-specific, so no global policy can be right. - pm.sample now only enforces sampler-independent constraints (step/nuts_sampler clashes, model identity, trace/callback/ return_inferencedata rejections) and forwards everything else verbatim: init, jitter_max_retries, discard_tuned_samples, keep_warning_stat, compile_kwargs. - ExternalSampler.sample defines the forwarded set as a closed contract (no **kwargs): a newly forwarded argument breaks every sampler loudly, forcing an explicit disposition. - Blackjax.sample documents and implements its dispositions with the reasoning for each: init is reduced to its jitter/no-jitter part (blackjax runs its own adaptation; jitter_initial_points overrides directly); a non-jax compile mode/backend raises instead of being ignored (unlike nutpie, blackjax can only consume a jax logp); discard_tuned_samples=False and keep_warning_stat warn with reasons. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
@ricardoV94 You're right that the previous iteration imposed NUTS semantics globally — a shared disposition table in
Each sampler owns its dispositions, with the argument for each.
Each disposition has a test ( 🤖 Generated with Claude Code |
|
@ricardoV94 I share your frustration, let me know if it got it this time. If not perhaps we'll scratch it. |
|
You won't mind my bot reply? Draft reply for #8353This iteration is better. Dropping the shared disposition table and making each sampler own its arguments fixed what I complained about last round. But I still don't want to land the The root issue is that pm.sample is married to NUTS. It has arguments ( A note on naming before the options: I renamed Options as I see them: Option 1: sampler= with configured objects (the current PR shape)pm.sample(sampler=pm.blackjax.nuts(target_accept=0.9), jitter_max_retries=10)The user has to split arguments between the sampler constructor and pm.sample, and which argument goes where is a historical accident of what pm.sample happened to list explicitly. It also keeps two classes of samplers: NUTS configured through pm.sample, everyone else through constructor objects. And pm.sample keeps having to police clashes per sampler. It also keeps the disposition problem from above: pm.sample keeps promising NUTS arguments like Option 2: same shape, but all NUTS arguments demoted out of pm.sampleConstructor takes algorithm configuration, pm.sample takes run configuration: pm.sample(sampler=pm.blackjax.nuts(target_accept=0.9), draws=1000, chains=4, random_seed=1)This is more defensible than option 1 and I don't think it's terrible. My hesitation: the object is instantiated with half the information it needs and then sits there waiting for pm.sample to hand it the other half (model, draws, chains, seeds) before it can do any real work. It's a configuration holder. The test is the criterion above: The rule also still forces a shared run vocabulary on everyone: every sampler must accept Option 3: pm.sample as a thin dispatcherOne flat kwargs namespace per sampler, and pm.sample forwards everything verbatim (the sampler is passed unconfigured, since its configuration is just the kwargs): pm.sample(sampler=pm.blackjax.mclmc, draws=1000, step_size=0.1)This fixes the splitting but hurts discoverability: the signature lives on one class while you type the arguments into another function. Option 4: no sampler argument in pm.sample at allEach sampler exposes its own pm.sample(draws=1000, ...) # the default/auto sampler
pm.nuts.sample(...)
pm.nutpie.nuts.sample(...) # Or: if a library only offers one sampler, it could be pm.library.sample()
pm.blackjax.mclmc.sample(model=m, draws=1000, step_size=0.1)
pm.numpyro.nuts.sample(model=m, draws=1000, target_accept=0.9)
pm.smc.sample(...)This is the only option with a single flat call site: options 1-3 all spread the arguments over two places in some form, and the funnel is what forces sampler objects into existence in the first place. Standalone flat sample functions don't need them. It's also somewhat orthogonal to 1-3. Many of these already exist ( To avoid N drifting copies of the run-level signature: the run-level arguments (draws, tune, chains, seeds, progressbar, idata_kwargs, ...) should live in one shared place, a common base class or run-options struct that every sample entry point takes, with a test asserting the signatures stay in sync. The internal sampler uses the same one, which is also what makes the demotion criterion testable. Same idea one level down: samplers of the same family (all the NUTS implementations, say) should share an abstract base class so the same argument doesn't drift into three spellings across libraries. What about step=?step= should be subsumable by the same interface, even if we keep the old argument as a shorthand: the internal step-method machinery (step assignment, CompoundStep, the python sampling loop) is itself a sampler, so something like This also gives The obstacle is just refactoring: native step samplers today eagerly create a lot of things at construction (model binding, compiled functions), so they can't act as pure configuration holders. They need a config phase separated from a build(model) phase. The same applies to this PR's What I think is best2 + 4. Every sampler gets a standalone flat For this PR concretely: the blackjax driver internals (kwargs routing, adaptation table, stats flattening, shared idata postprocessing) are good and worth keeping as they are. I'd have it ship both surfaces already: the flat To make the end state concrete, pm.sample after all the deprecations are through: pm.sample(
# what to sample
draws=1000,
tune=1000,
chains=None,
cores=None,
model=None, # taken from context if not given
# who samples it and how it's compiled
sampler=None, # None resolves a default
backend=None, # also steers the default sampler choice
compile_kwargs=None,
# how the run goes
initvals=None,
random_seed=None,
progressbar=True, # not eagerly a bool, nutpie already handles more refined specs
# what comes back
discard_tuned_samples=True,
keep_warning_stat=False,
var_names=None,
idata_kwargs=None,
compute_convergence_checks=True,
)Everything left is run configuration that any sampler in the funnel understands ( And the shape of the classes behind it: class Sampler(ABC):
"""Base for everything pm.sample(sampler=...) accepts.
The constructor is algorithm configuration only: nothing model-bound,
nothing compiled. Model and run configuration arrive at sample time.
"""
@abstractmethod
def sample_from_init(
self,
*,
model,
draws,
tune,
chains,
cores,
initvals,
random_seed,
progressbar,
discard_tuned_samples,
keep_warning_stat,
var_names,
idata_kwargs,
compute_convergence_checks,
compile_kwargs,
) -> InferenceData:
"""Exactly pm.sample's signature minus sampler/backend (enforced by a test).
Not called `sample` so it can't be confused with the flat entry points.
"""
class ExternalSampler(Sampler):
"""Shared implementation logic for samplers living in other libraries."""
package: str # checked at construction, raises with an install hint
version_range: str | None = None # optional, e.g. ">=1.2,<2"
def __init__(self):
require_installed(self.package, self.version_range)
class Blackjax(ExternalSampler):
package = "blackjax"
def __init__(self, algorithm="nuts", adaptation=None, jitter_initial_points=True, **algorithm_kwargs):
# resolve algorithm in the blackjax namespace, route kwargs by
# signature introspection, raise on unknown or missing ones.
# No model, nothing compiled.
...
def sample_from_init(self, *, model, draws, ...):
# jaxify the model logp, adaptation + sampling loop, return idata
...
class AlgorithmEntry:
"""What pm.blackjax.nuts actually is: both surfaces in one object."""
def __call__(self, **algorithm_kwargs) -> Blackjax:
# option 2 surface: pm.blackjax.nuts(target_accept=0.9)
return Blackjax(self.algorithm, **algorithm_kwargs)
def sample(self, *, model=None, draws=1000, ..., **algorithm_kwargs) -> InferenceData:
# option 4 surface: pm.blackjax.nuts.sample(...), one flat call
return self(**algorithm_kwargs).sample_from_init(model=model, draws=draws, ...)
class StepSampler(Sampler):
"""The native machinery: step assignment, CompoundStep, the python loop."""
def __init__(self, step=None): # None: automatic assignment at sample time
...Note where the internal/external distinction ended up: it's implementation inheritance, not interface. |
Implements the design agreed in review: every sampler exposes a flat entry point with a full explicit signature, and pm.sample keeps a sampler= funnel as sugar over it, under the rule "constructor takes algorithm configuration, pm.sample takes run configuration". - Sampler ABC: constructor is algorithm configuration only (nothing model-bound, nothing compiled); the model and run configuration arrive at sample time via sample_from_init, whose signature is the shared run contract (kept in sync with pm.sample by a test). The method is deliberately not named `sample` so it cannot be confused with the flat entry points. - ExternalSampler is implementation inheritance only: a dependency check (package/version_range) with an install hint, nothing added to the Sampler contract. - pm.blackjax.<algo> is an AlgorithmEntry exposing both surfaces: calling it configures a Blackjax sampler for the funnel; .sample(...) draws in one flat call, defined in terms of the first so the two cannot disagree. - pm.sample(sampler=...) forwards only run configuration and now rejects the NUTS-specific arguments (init, n_init, jitter_max_retries) instead of forwarding them for reinterpretation; jitter configuration moved into the Blackjax constructor (jitter_initial_points, jitter_max_retries). Demoting the NUTS arguments out of pm.sample for the internal sampler is left for a follow-up. - Renames: external_sampler= -> sampler=, pm.external.blackjax -> pm.blackjax, pymc/sampling/external/ -> pymc/sampling/samplers/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Implemented your 2+4 design in d6345fd, mapping directly onto the sketch: Both surfaces, one implementation. pm.sample(sampler=pm.blackjax.nuts(target_accept=0.9))
pm.blackjax.mclmc.sample(model=m, draws=1000, step_size=0.1)The class shape. Constructor is pure configuration. No The funnel rejects NUTS vocabulary instead of forwarding it. With Renames throughout: Left for follow-ups, per your scoping: demoting 🤖 Generated with Claude Code |
|
can you open a PR where you apply the refactor to the existing external samplers, and the step_sampler, but don't add all the new blackjax support? I'd like to review/merge that as a unit first. Add a DeprecationWarning for nuts_sampler (not FutureWarning yet to give a bit more room to transition). And for the NUTS - specific arguments. Basically refactor all the existing code to the new format. Then we can add the new blackjax samplers. Otherwise I fear code rot, as there won't be any interest in cleaning up when the functionality is already there |
pyarrow 25.0.0, pulled in transitively when tests.yml later installs nutpie via pip, segfaults XLA CPU compilation when both are loaded in the same process. nutpie only requires pyarrow>=12.0.0 with no upper bound, so pre-installing a pinned pyarrow here via conda/pip satisfies that requirement and pip leaves it alone during the nutpie install. This achieves the same effect as pinning directly in the workflow's nutpie install step (see PR pymc-devs#8353), without touching .github/workflows/tests.yml.
|
As requested: the refactor-only unit is split out into #8358 (existing samplers + StepSampler behind the shared Sampler interface, nuts_sampler/init deprecations, no new blackjax support). Once that merges I'll rebase this PR down to just the generic blackjax driver on top of it. |
2 participants
What
A generic external sampler that can run any blackjax MCMC algorithm on a PyMC model, following up on the design discussion in #7880 (comment) and #7699:
How
blackjaxnamespace and are validated viaisinstance(obj, blackjax.GenerateSamplingAPI)— new blackjax algorithms (slice sampling, nested sampling, GIST in 1.6) will work with zero changes here. SMC/SGMCMC/VI and cross-chain adaptation raise informative errors pointing at alternatives.Blackjax(...).get_kwargs()shows every accepted parameter with its default; unknown kwargs and missing required parameters (e.g.Blackjax("mala")withoutstep_size) raise at construction time with the full accepted-parameter listing. This addresses thenuts_sampler_kwargsdiscoverability problem.window(alsolow_rank/pathfinder) for nuts/hmc/dynamic_hmc,mclmc_find_L_and_step_sizefor MCLMC,adaptation=None+ plain burn-in for static-parameter algorithms (mala, barker, rmh, ghmc, ...). User-overridable viaadaptation=.Infonamedtuple is flattened generically (scalars only) and renamed to ArviZ conventions (diverging,tree_depth,n_steps,lp, ...), so every algorithm gets sensiblesample_statswithout per-algorithm code.sample_jax_nutsis extracted into a shared_postprocess_and_build_idatahelper (pure refactor, existing numpyro/blackjax NUTS behavior unchanged).On the
pm.sampleside this adds only a minimalexternal_sampler=hook (mirroring #7880): validation of clashes (step,nuts_sampler,trace,callback, leftover NUTS kwargs → error directing users to configure the sampler at construction) and delegation. The existingnuts_sampler=path is untouched, so this PR does not depend on resolving the argument-migration questions in #7880 — it can rebase onto it when that lands.Scope
Phase 1 = plain MCMC + single-chain adaptation. Deliberately excluded (informative errors): SMC (needs split logprior/loglikelihood; incubating in pymc-extras), SGMCMC (needs minibatch plumbing), MEADS/ChEES (cross-chain adaptation, needs a second driver — can be a follow-up), VI (fit-then-sample shaped, belongs next to
pymc_extras.fit).@ricardoV94 re "not sure if we want 15 string versions into external.blackjax": the class accepts either the string or the blackjax object itself (
Blackjax(blackjax.mclmc)), so the string is never load-bearing — happy to add per-algorithm aliases (pm.external.blackjax.MCLMC()) on top if preferred; that's a thin layer over this driver either way.Tests
tests/sampling/test_external.py: parameter recovery for nuts/hmc/mclmc, runs for mala/barker, kwargs routing, stat renames, and the full error matrix (unknown algorithm/kwarg, missing required params, VI/SMC/SGMCMC rejection, discrete models,pm.sampleclashes, model mismatch).🤖 Generated with Claude Code