BlackJAX: JAX samplers split into composable transition kernels

It is a sampler library, not a probabilistic programming language

The README is explicit about what BlackJAX is not, and it opens on that point rather than burying it: BlackJAX is a library of samplers for JAX that works on CPU as well as GPU, and it is not a probabilistic programming library. What it requires from you is a log probability density function, potentially unnormalized, that is compatible with JAX. The library describes itself as bridging the gap between one liner frameworks and modular, customizable libraries, and the two audiences it names are people who have a logpdf and just need a sampler, and developers of probabilistic programming languages who need samplers they can call. If your work is writing the model, this library is downstream of that work, and you would be providing the density function yourself.

The whole design is one kernel signature: state in, state out

Everything in BlackJAX is built around a single pattern, and the README states it in one line of code. Anything that takes a state and returns a state is a transition kernel, implemented as new_state, info = kernel(rng_key, state). Kernels are stateless functions and all follow the same API, with the state and the information about the transition returned separately, which is what makes them easy to compose and exchange. Specialisation happens by closure rather than by passing parameters, so a kernel with a fixed step size is a different object from a kernel that takes one as an argument. The design argument in the README is that sampling algorithms are too often integrated into probabilistic programming libraries and not decoupled from the rest of the framework, which makes them hard to use for people who do not need the modelling language, and that the implementation is most of the time monolithic so parts cannot be reused to build custom kernels. Both problems are solved by making the kernel the unit of composition.

The building blocks are integrators, proposals and momentum generators

The claimed strength of the library is in its internals rather than its surface, and the specific components named are the integrators, the proposal, the momentum generators and so on. The point of exposing them is that you can combine them to build new algorithms, and the README frames this as an opportunity to accelerate research on sampling algorithms by providing reusable code. New algorithms are meant to be tested by reusing the basic components of methods already implemented, following the developer guidelines, so that a new idea can be tried against existing methods without writing everything from scratch. In practice this is the difference between a library that gives you NUTS and a library that lets you write your own Hamiltonian Monte Carlo variant in an afternoon by swapping the integrator while keeping the momentum distribution. The README also notes that the NUTS implementation was largely inspired by Numpyro.

Installing on CPU by default, and the GPU trap in the quickstart

Installation is a pip or conda-forge command, but the more important detail is the caveat that follows. BlackJAX is written in pure Python but depends on XLA via JAX, and by default the version of JAX installed alongside it will make your code run on CPU only. If you want GPU or TPU, the README tells you to follow the JAX installation instructions to install JAX with the relevant hardware acceleration support, which means the accelerator setup is on you and not something the package handles. The packaging metadata underlines the floor: the project requires Python 3.13 or newer, declares jax and jaxlib at 0.9.0 or above, numpy 1.25, optax 0.2.3, scipy 1.13 and typing-extensions 4.4.0. There is also an optional extra called progress, which adds tqdm and jax-tap.

pip install blackjax

conda install -c conda-forge blackjax

A first NUTS run: build the kernel, init the state, iterate

The quickstart example is short and it shows the API shape better than any description. You define a logdensity_fn over a dictionary of parameters, here a location and a scale for a normal distribution, build the kernel with blackjax.nuts by passing the density function, a step size and an inverse mass matrix, initialise the state with nuts.init on an initial position, and then iterate. The loop folds a new key per iteration with jax.random.fold_in and calls the jitted step function, so the sampler is a pure function of the key and the state. Two details are worth copying exactly. The step function is wrapped in jax.jit, so compilation happens once rather than per iteration, and the key is folded per iteration rather than reused, which is what keeps the chain reproducible. The README also links to documentation covering inference loops for one or several chains and the use of the Stan warmup.

import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
import numpy as np

import blackjax

observed = np.random.normal(10, 20, size=1_000)
def logdensity_fn(x):
    logpdf = stats.norm.logpdf(observed, x["loc"], x["scale"])
    return jnp.sum(logpdf)

step_size = 1e-3
inverse_mass_matrix = jnp.array([1., 1.])
nuts = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)

initial_position = {"loc": 1., "scale": 2.}
state = nuts.init(initial_position)

rng_key = jax.random.key(0)
step = jax.jit(nuts.step)
for i in range(100):
    nuts_key = jax.random.fold_in(rng_key, i)
    state, _ = step(nuts_key, state)

The test suite runs on CPU with the race of parallel pytest workers

The Makefile shows how the project is tested and it makes two choices you would want to know about. The test target syncs the development group and the progress extra with uv, then runs pytest with the xdist parallel flag, verbose output, benchmarks disabled, and coverage reported both as XML and to the terminal, with the JAX platform pinned to cpu. Pinning the platform to cpu for tests means the suite does not require an accelerator, and the parallel execution means the suite is written to be safe when many workers run at once, which is a good sign for a library whose kernels are meant to be pure functions. The release target is equally simple: the version number is read from git with setuptools_scm, and releasing means tagging and pushing the tag, so the version in the published package is derived from the commit history rather than written by hand.

Governance, licence and release cadence

The repository carries a GOVERNANCE.md, a CODE_OF_CONDUCT.md and a CONTRIBUTING.md, so this is a project with an explicit governance structure rather than a single maintainer's repository, and the README asks contributors to follow a short contribution guide. The licence is Apache 2.0, with the text in the LICENSE file, and the packaging metadata classifies the development status as Production/Stable. The citation is a BibTeX entry for the paper BlackJAX: Composable Bayesian inference in JAX by Cabezas, Corenflos, Lao and Louf, with an arXiv eprint identifier, and the README notes that in that entry the version number should be the last tag on the main branch. On upkeep, the last push to the repository was on 2026-10-08, with releases 1.7.0 and 1.7.1 both on 2026-10-02 and 1.6.2 on 2026-07-16, so two version bumps landing on the same day is normal here and a dependency on BlackJAX should be checked against a specific tag rather than a branch.

Editorial conclusion

BlackJAX suits a researcher or library author who already has a log probability function in JAX and wants a sampler, or who wants to assemble a new MCMC method from integrators, proposal distributions and momentum generators rather than fork a monolithic implementation. It is the wrong choice if you want a modelling language, if your log density is not expressible in JAX, and if you are on a Python version below 3.13, since the packaging metadata requires 3.13 as a floor. Before adopting it, confirm the jax and jaxlib version your environment already has meets the 0.9.0 floor the package declares, because a plain pip install pulls a CPU only JAX unless you install the accelerated build yourself, and read the developer guidelines for how a new kernel is expected to be structured and tested.

Frequently asked questions

Official sources

1. blackjax-devs/blackjax on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases