Registry / ai-ml / blackjax

blackjax

JSON →
library1.5pypypi✓ verified 85d ago

Blackjax is a flexible and fast Markov chain Monte Carlo (MCMC) sampling library in Python, built on JAX for GPU/TPU acceleration. Current version is 1.5, with active development and frequent releases.

pip install blackjax
INSTALL
IMPORT
SIG · BLACKJAX
B
blackjax
ai-mlpythonv1.5
Install
13.6s avg
Import
4323ms
Disk
337MB
Pass rate
10/ 10
Env Coverage10 / 10
glibc
3.93.13
musl
3.93.13
Install & Compatibility
Where this runs
tested against v0.8.2 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
musl
py 3.103.95 runs
installs and imports cleanly · install 0.0s · import 0.000s · 56.7MB
glibc
py 3.103.95 runs
installs and imports cleanly · install 13.6s · import 3.458s · 638MB
337MB installed
● package 337MB
Code
Verified usage

Verified import paths — ran on the pinned version, not inferred.

blackjax
import blackjax
from blackjax import ... (not for top-level)
Top-level package; submodules like blackjax.mcmc are accessed via dot notation.

Minimal HMC sampling with Blackjax.

import jax import jax.numpy as jnp import blackjax # Define a simple target distribution (2D Gaussian) def logdensity_fn(x): return -0.5 * jnp.sum(x**2) # Build the HMC kernel kernel = blackjax.hmc(logdensity_fn, step_size=0.1, inverse_mass_matrix=jnp.eye(2), num_integration_steps=10) # Initialize state key = jax.random.PRNGKey(0) initial_position = jnp.array([1.0, 1.0]) initial_state = kernel.init(initial_position) # Sample key, subkey = jax.random.split(key) state, info = kernel.step(subkey, initial_state) print(state.position)
Debug
Known issues
breakingBlackjax v1.0+ removed the old API using `blackjax.mcmc` sampler constructors (e.g., `blackjax.mcmc.hmc`). Use `blackjax.hmc` directly.
fix
Replace `from blackjax.mcmc import hmc` with `import blackjax; kernel = blackjax.hmc(...)`.
affects: <1.0
gotchaBlackjax requires JAX and does not support plain NumPy arrays as inputs; all functions must use JAX numpy (`jnp`) and PRNG keys.
fix
Always use `jax.random.PRNGKey` for randomness and `jnp.array` for data.
affects: All
deprecatedThe `blackjax.mcmc` submodule is deprecated in favor of top-level sampler functions (e.g., `blackjax.hmc`, `blackjax.nuts`).
fix
Use `blackjax.hmc(...)` instead of `blackjax.mcmc.hmc(...)`.
affects: >=1.0
Errors
Common errors & fixes
ModuleNotFoundError: No module named 'blackjax'
Blackjax not installed or installed in an isolated environment.
fix
Run `pip install blackjax` in the correct environment.
AttributeError: module 'blackjax' has no attribute 'mcmc'
Using old import path for samplers.
fix
Directly call `blackjax.hmc(...)` or `blackjax.nuts(...)` instead of `blackjax.mcmc.hmc(...)`.
Upgrade
Version history
1.5latest on PyPI · released Apr 3, 2026
Audit
Dependencies
jaxrequiredCore dependency for computation and autograd
jaxlibrequiredJAX library for CPU/GPU/TPU backends
jaxtypingoptionalType annotations for JAX arrays
Agent activity
17 hits · last 30 days
node
12
Amazon
1
OpenAI (training)
1
Resources
blackjax — pip install blackjax · libregistry