blackjax

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 on this page — 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

blackjax on PyPI

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

Evidence: blackjax-1.6.2-py3-none-any.whl

Keywords: probability, machine learning, statistics, mcmc, sampling

Tags

More Scientific/Engineering packages