# BlackJAX: Composable Bayesian inference in JAX

arXiv: https://arxiv.org/abs/2402.10797 · HTML: https://mlsanity.com/paper/2402.10797/

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 45 references that are in the map.
- First posted on arXiv 16 Feb 2024.

## Builds on (most cited first, 9 of 9)

- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [Bayesian Layers: A Module for Neural Network Uncertainty](https://mlsanity.com/paper/1812.03973/) (2018)
- [Bayes-Newton Methods for Approximate Bayesian Inference with PSD Guarantees](https://mlsanity.com/paper/2111.01721/) (2021)
- [Transport Elliptical Slice Sampling](https://mlsanity.com/paper/2210.10644/) (2022)
- [Identifying the Optimal Integration Time in Hamiltonian Monte Carlo](https://mlsanity.com/paper/1601.00225/) (2016)
- [Non-reversibly updating a uniform [0,1] value for Metropolis accept/reject decisions](https://mlsanity.com/paper/2001.11950/) (2020)
- [Benchmarking Bayesian neural networks and evaluation metrics for regression tasks](https://mlsanity.com/paper/2206.06779/) (2022)
- [Bayesian cross-validation by parallel Markov Chain Monte Carlo](https://mlsanity.com/paper/2310.07002/) (2023)

## Built on by (most cited first, 26 of 26)

- [Sequential Controlled Langevin Diffusions](https://mlsanity.com/paper/2412.07081/) (2024)
- [BayesFlow 2: Multi-Backend Amortized Bayesian Inference in Python](https://mlsanity.com/paper/2602.07098/) (2026)
- [Simulation-based Inference with the Python Package sbijax](https://mlsanity.com/paper/2409.19435/) (2024)
- [On the Interplay of Priors and Overparametrization in Bayesian Neural Network Posteriors](https://mlsanity.com/paper/2603.22030/) (2026)
- [Uncertainty in Physics and AI: Taxonomy, Quantification, and Validation](https://mlsanity.com/paper/2605.10378/) (2026)
- [Geodesic Slice Sampler for Multimodal Distributions with Strong Curvature](https://mlsanity.com/paper/2502.21190/) (2025)
- [Fisher meets Feynman: score-based variational inference with a product of experts](https://mlsanity.com/paper/2510.21598/) (2025)
- [Statistical Learning Analysis of Physics-Informed Neural Networks](https://mlsanity.com/paper/2602.11097/) (2026)
- [How to Approximate Inference with Subtractive Mixture Models](https://mlsanity.com/paper/2604.16714/) (2026)
- [bde: A Python Package for Bayesian Deep Ensembles via MILE](https://mlsanity.com/paper/2605.14146/) (2026)
- [Federated Learning for Non-factorizable Models using Deep Generative Prior Approximations](https://mlsanity.com/paper/2405.16055/) (2024)
- [Diffusion Generative Modelling for Divide-and-Conquer MCMC](https://mlsanity.com/paper/2406.11664/) (2024)
- [Scalable Monte Carlo for Bayesian Learning](https://mlsanity.com/paper/2407.12751/) (2024)
- [Flow Annealed Importance Sampling Bootstrap meets Differentiable Particle Physics](https://mlsanity.com/paper/2411.16234/) (2024)
- [Efficiently Vectorized MCMC on Modern Accelerators](https://mlsanity.com/paper/2503.17405/) (2025)
- [Bayesian Inverse Physics for Neuro-Symbolic Robot Learning](https://mlsanity.com/paper/2506.08756/) (2025)
- [Reducing normalizing flow complexity for MCMC preconditioning](https://mlsanity.com/paper/2511.02345/) (2025)
- [CLAX: Fast and Flexible Neural Click Models in JAX](https://mlsanity.com/paper/2511.03620/) (2025)
- [Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics](https://mlsanity.com/paper/2604.21456/) (2026)
- [Corrected Integrated Laplace Approximation for Bayesian Inference in Latent Gaussian Models](https://mlsanity.com/paper/2605.20345/) (2026)
- [GenSBI: Generative Methods for Simulation-Based Inference in JAX](https://mlsanity.com/paper/2605.27499/) (2026)
- [Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms](https://mlsanity.com/paper/2607.24960/) (2026)
- [Guiding Posterior Exploration with Optimizer-Derived Geometry](https://mlsanity.com/paper/2607.25312/) (2026)
- [Subspace Inference Enables Efficient Active Reward Learning from Preferences](https://mlsanity.com/paper/2609.04066/) (2026)
- [Multiple latent orderings better predict language model preferences](https://mlsanity.com/paper/2609.22170/) (2026)
- [Predictively Oriented Gaussian Process Posteriors](https://mlsanity.com/paper/2610.03201/) (2026)

## Abstract (arXiv)

BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by taking a functional approach to the algorithms' implementation. BlackJAX is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs. The library integrates well with probabilistic programming languages by working directly with the (un-normalized) target log density function. BlackJAX is intended as a collection of low-level, composable implementations of basic statistical 'atoms' that can be combined to perform well-defined Bayesian inference, but also provides high-level routines for ease of use. It is designed for users who need cutting-edge methods, researchers who want to create complex sampling methods, and people who want to learn how these work.