Bayesian sampling, decomposed: a JAX-native sampler toolkit

What it does BlackJAX is a library of Bayesian sampling algorithms—NUTS, HMC, and friends—built on JAX. It runs on CPU or GPU and targets developers who already have a log-probability density function and just need a fast, reliable sampler. It is explicitly not a probabilistic programming language, though it plays nicely with PPLs that can export JAX-compatible logpdfs.

The interesting bit The whole library is organized around a single, stubbornly simple contract: every transition kernel is a stateless function new_state, info = kernel(rng_key, state). That uniformity lets you swap integrators, proposals, and momentum generators like Lego bricks. The authors call this “specialization by closure instead of passing parameters”—a fancy way of saying you build custom samplers by wiring together small, testable pieces rather than forking monolithic codebases.

Key highlights

• Runs on CPU, GPU, and TPU via JAX/XLA
• NUTS implementation draws from NumPyro’s battle-tested approach
• Lower-level building blocks (integrators, proposals, etc.) exposed for research use
• Stateless kernels with a uniform API make composition and testing straightforward
• Pure Python, though performance depends on JAX’s compilation pipeline

Caveats

• Requires manual JAX installation with hardware-specific backends for GPU/TPU use; the default pip install is CPU-only
• Not a modeling language—you bring your own logpdf, which means no automatic differentiation of graphical models

Verdict Grab this if you’re building a PPL, hacking on novel sampling algorithms, or just need NUTS on a GPU without the kitchen sink of a full probabilistic framework. Skip it if you want a point-and-click Bayesian workflow or need the modeling layer to hold your hand.