---
id: blackjax
version: "1.6.2"
license: Apache License 2.0
license_treatment: permissive
maintenance: active
---

# blackjax — Flexible and fast sampling in Python

License: permissive · Maintenance: active · Downloads: 227.1K/mo

## What it is and what it does

BlackJAX is a library of sampling algorithms—primarily MCMC methods like NUTS—built on top of JAX and designed to run efficiently on both CPU and GPU hardware. It is not a probabilistic programming language itself, but rather a collection of composable, reusable sampling kernels that integrate well with any tool that can provide a log-probability density function compatible with JAX.

The library targets two audiences: users who have a logpdf and just need a working sampler (via simple high-level APIs like `blackjax.nuts()`), and researchers building custom inference algorithms who want to reuse robust, well-tested building blocks like integrators, proposal generators, and momentum samplers. All kernels follow a stateless functional pattern—taking a random key and state, returning a new state and metadata—making them easy to compose and swap. The package depends on jax, jaxlib, numpy, optax, scipy, and typing-extensions.

Use it for:

- Sample from a posterior distribution when you have a log-probability function but no full probabilistic programming framework.
- Build custom MCMC samplers by composing elementary kernels (integrators, proposals, momentum generators) without reimplementing from scratch.
- Run inference on GPU or TPU by leveraging JAX's hardware acceleration for large-scale Bayesian computations.
- Integrate sampling into an existing probabilistic programming language by providing a modular, decoupled sampler.
- Learn how sampling algorithms work by studying and experimenting with the library's composable, well-documented building blocks.
- Accelerate research on new sampling schemes using robust, performant, and reusable algorithm components.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

BlackJAX provides a collection of MCMC and other sampling algorithms built on JAX that run on CPU and GPU, designed for users who have a log-probability function and need a performant sampler.

Yes. BlackJAX is actively maintained, has no known vulnerabilities, installs with low friction, and is licensed permissively. It is well-suited if you need MCMC sampling on JAX with GPU support, want composable algorithm building blocks, or are integrating sampling into a larger system. Not necessary if you only need sampling from standard distributions or are already using a full probabilistic programming framework.

## Install

pip install blackjax
uv add blackjax
poetry add blackjax

## Installing blackjax

Before you install: Installation is straightforward via pip with low friction. The package is actively maintained with a recent release (29 days ago) and receives regular updates. Note that JAX defaults to CPU-only; GPU/TPU support requires separate JAX installation following JAX's hardware-specific instructions.

License in practice: Licensed under Apache License 2.0 (permissive), allowing use in commercial and proprietary projects with minimal restrictions—you must include a copy of the license and state significant changes, but can use the code freely.

Quickstart:

pip install blackjax

import blackjax
import jax
import jax.numpy as jnp

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

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

rng_key = jax.random.key(0)
state, info = nuts.step(rng_key, state)

Requires Python 3.11 or later. GPU/TPU sampling requires separate JAX installation with hardware support; default JAX install runs on CPU only.

Verify before relying:

- Whether the library supports all modern JAX features and remains compatible with future JAX releases.
- Performance benchmarks comparing BlackJAX samplers to other MCMC libraries on standard problems.

## Package facts

- License: Apache License 2.0 (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 227.1K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

mcmc sampling jax, bayesian inference gpu, nuts sampler python, probabilistic sampling library, jax-based mcmc, hamiltonian monte carlo, composable sampling algorithms, mcmc, bayesian-inference, gpu-accelerated

[View on SkillFed](https://skillfed.io/packages/blackjax) · [View on PyPI](https://pypi.org/project/blackjax/)