BlackJAX: Composable Bayesian inference in JAX
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.
Cite
@article{arxiv.2402.10797,
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},
journal= {arXiv preprint arXiv:2402.10797},
year = {2024}
}
Comments
Companion paper for the library https://github.com/blackjax-devs/blackjax Update: minor changes and updated the list of authors to include technical contributors