BlackJAX is a library of samplers for JAX that runs on CPU, GPU, and TPU. It is not a probabilistic programming library — but it integrates seamlessly with any PPL that can expose a JAX-compatible log-probability density function.

It bridges the gap between one-liner frameworks and fully modular, composable inference toolkits. Use it as a black box, or crack it open and compose your own algorithms from low-level building blocks.

BlackJAX ships a broad collection of production-ready samplers:

pip install blackjax

import jax
import jax.numpy as jnp
import blackjax

def logdensity_fn(x):
    return -0.5 * jnp.sum(x ** 2)  # standard normal

# Build a NUTS sampler
nuts = blackjax.nuts(logdensity_fn, step_size=1e-3, inverse_mass_matrix=jnp.ones(2))

# Initialize and run a single step
state = nuts.init(jnp.zeros(2))
rng_key = jax.random.key(0)
state, info = jax.jit(nuts.step)(rng_key, state)

See the documentation for full inference loops, multi-chain sampling, Stan warmup, and more.

BlackJAX is built by researchers and engineers who care about sampling. We welcome contributions of all kinds — new algorithms, bug fixes, documentation, and discussions.

• 🛠️ Contribution guidelines — how to get started
• 🏛️ Governance — how the project is run
• 💬 Open an issue — bug reports, feature requests, or questions
• 🌱 Developer guidelines — implement and test a new sampler