BlackJAX: Composable Bayesian inference in JAX

TL;DR

BlackJAX is a Python library built on JAX that offers modular, high-performance Bayesian inference algorithms, enabling flexible and efficient sampling and variational methods across various hardware platforms.

Contribution

It introduces a composable, low-level implementation of Bayesian inference algorithms in JAX, facilitating customization and high performance for research and practical applications.

Findings

Supports CPUs, GPUs, and TPUs with JAX compilation

Provides both low-level primitives and high-level routines

Enables flexible combination of inference components

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.