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

---

**What changed vs. the current 3-line README:**

| Before | After |
|---|---|
| No visual identity | Centered hero with the existing scatter.gif animation |
| No badges | CI, codecov, PyPI, arXiv, stars badges |
| No description | Clear pitch explaining what BlackJAX is and who it's for |
| No algorithm listing | Table covering MCMC, SMC, SG, variational, orbit families |
| No code | Minimal NUTS quick-start snippet |
| 2 bare links | Rich resource table + community section |