BlackJAX: Composable Bayesian inference in JAX
This provides a modular and fast library for users, researchers, and learners in Bayesian computation, though it is incremental as it builds on existing algorithms.
The authors tackled the problem of implementing Bayesian inference algorithms by developing BlackJAX, a library that provides composable and efficient sampling and variational methods in JAX, resulting in a tool that runs on CPUs, GPUs, and TPUs and integrates well with probabilistic programming languages.
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.