blackjax-devs/blackjax

Health v2 · maintenance only

Adoption not part of the score

1119 stars · 150 forks observed · 2026-09-01

What it is AI-extracted, prompt v1, taxonomy v1, 2026-08-30, confidence not recorded

BlackJAX is a library of fast, modular MCMC samplers (NUTS, HMC, and more) built on JAX, running on CPU, GPU, and TPU. It is not a full probabilistic programming language but integrates with any PPL that can supply a JAX-compatible log-probability density function.

Use cases

• sample from a posterior given a logpdf
• run MCMC sampling on GPU
• build a custom sampling algorithm from modular kernels
• use NUTS or HMC with a probabilistic programming language
• learn how MCMC sampling algorithms work internally
• speed up Bayesian inference with JIT-compiled samplers

When to choose

• you already have a JAX-compatible log-density and just need a sampler
• you need GPU/TPU-accelerated sampling
• you are a researcher building or studying new sampling algorithms
• you are developing a PPL and need robust sampler building blocks

When to avoid

• you want a full probabilistic programming language with model syntax and inference automation
• you need a one-liner black-box inference framework like PyMC or Stan
• your project does not use JAX or cannot provide a JAX-compatible logpdf

Facets

library · maturity active

machine-learning simulation math machine-learning data-science python cross-platform mcmc bayesian-inference jax sampling hamiltonian-monte-carlo nuts probabilistic-programming statistics algorithms gpu

Member repositories

For agents

markdown · JSON · MCP: product_card(name="blackjax-devs/blackjax")