BlackJAX: Composable Bayesian inference in JAX

Alberto Cabezas, Adrien Corenflos, Junpeng Lao, Rémi Louf, Antoine Carnec +17 more
2/16/2024
cs.MScs.LGstat.CO

No official code found. Best implementation: blackjax-devs/blackjax (1,129 stars, Apache-2.0, last commit 2026-04-28).

Abstract

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.

View on arXivView PDF

Code Implementations(1)

blackjax-devs/blackjaxCommunity100%
Apache-2.0

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

1,129155Python, MakefileDec 9, 20205 months agoApache-2.0
bayesian-inferencehamiltonian-monte-carloprobabilistic-programmingsampling-methods

Similar Papers

Simulation-based inference with the Python Package sbijax

77% match1 repo552024

laplax -- Laplace Approximations with JAX

74% match2025

SynJax: Structured Probability Distributions for JAX

74% match2 repos2542023

BARK: A Fully Bayesian Tree Kernel for Black-box Optimization

74% match2025

GaussianProcesses.jl: A Nonparametric Bayes package for the Julia Language

72% match1 repo3202018

On the Convergence of Black-Box Variational Inference

71% match2023

Cite this paper

@article{cabezas2024blackjax,
  title  = {BlackJAX: Composable Bayesian inference in JAX},
  author = {Alberto Cabezas and Adrien Corenflos and Junpeng Lao and Rémi Louf and Antoine Carnec and Kaustubh Chaudhari and Reuben Cohn-Gordon and Jeremie Coullon and Wei Deng and Sam Duffield and Gerardo Durán-Martín and Marcin Elantkowski and Dan Foreman-Mackey and Michele Gregori and Carlos Iguaran and Ravin Kumar and Martin Lysy and Kevin Murphy and Juan Camilo Orduz and Karm Patel and Xi Wang and Rob Zinkov},
  year   = {2024},
  eprint = {2402.10797},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2402.10797v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2402.10797.svg)](https://codeofpaper.com/paper/2402.10797)

Discussion