BlackJAX: Composable Bayesian inference in JAX

BlackJAX: Composable Bayesian inference in JAX

2024 | Alberto Cabezas, Adrien Corenflos, Junpeng Lao, Rémi Louf
BlackJAX is a Python library for Bayesian inference using JAX, offering sampling and variational inference algorithms. It is designed for ease of use, speed, and modularity, allowing users to compose algorithms. BlackJAX supports MCMC, SMC, and variational inference, and provides a functional API for building new algorithms. It integrates with probabilistic programming languages by working directly with the target log density function. The library includes basic components for statistical inference that can be combined to perform Bayesian inference. BlackJAX is intended for users needing cutting-edge methods, researchers developing complex sampling methods, and learners wanting to understand how these work. BlackJAX provides composable inferential building blocks, such as Metropolis-Hastings, Hamiltonian dynamics, and stochastic gradient utilities, which can be used in sequential Monte Carlo or mean field approximations. These components are unified under a convenient API that allows for the creation of new or existing algorithms. The library's functional design enables users to build and experiment with new algorithms by applying the same mathematical logic used to design them. BlackJAX supports a range of sampling algorithms, including MCMC, SMC, and Stochastic Gradient MCMC, as well as approximate inference algorithms like variational inference. It is designed to be efficient, with auxiliary information included in the state for performance. The library is compatible with several JAX-powered scientific libraries and is used in courses and tutorials. BlackJAX has contributed to the development of Bayesian inference methods and has been adopted in various research papers. It has also been used in educational materials, such as Kevin Murphy's authoritative manuscript. The future of BlackJAX includes expanding its Bayesian computation methods, adding performance diagnostic tools, and improving documentation and tutorials. The library is open-source, with contributions welcome from the community. It follows a self-appointing council model for governance and has a comprehensive test suite with 99% test coverage.BlackJAX is a Python library for Bayesian inference using JAX, offering sampling and variational inference algorithms. It is designed for ease of use, speed, and modularity, allowing users to compose algorithms. BlackJAX supports MCMC, SMC, and variational inference, and provides a functional API for building new algorithms. It integrates with probabilistic programming languages by working directly with the target log density function. The library includes basic components for statistical inference that can be combined to perform Bayesian inference. BlackJAX is intended for users needing cutting-edge methods, researchers developing complex sampling methods, and learners wanting to understand how these work. BlackJAX provides composable inferential building blocks, such as Metropolis-Hastings, Hamiltonian dynamics, and stochastic gradient utilities, which can be used in sequential Monte Carlo or mean field approximations. These components are unified under a convenient API that allows for the creation of new or existing algorithms. The library's functional design enables users to build and experiment with new algorithms by applying the same mathematical logic used to design them. BlackJAX supports a range of sampling algorithms, including MCMC, SMC, and Stochastic Gradient MCMC, as well as approximate inference algorithms like variational inference. It is designed to be efficient, with auxiliary information included in the state for performance. The library is compatible with several JAX-powered scientific libraries and is used in courses and tutorials. BlackJAX has contributed to the development of Bayesian inference methods and has been adopted in various research papers. It has also been used in educational materials, such as Kevin Murphy's authoritative manuscript. The future of BlackJAX includes expanding its Bayesian computation methods, adding performance diagnostic tools, and improving documentation and tutorials. The library is open-source, with contributions welcome from the community. It follows a self-appointing council model for governance and has a comprehensive test suite with 99% test coverage.
Reach us at info@study.space