BlackJAX: Composable Bayesian inference in JAX

2024/02/16 by Alberto Cabezas, Cabezas, Alberto, Adrien Corenflos +41 · 1 voice · 22 citations
Computer Science · Mathematics · #Bayesian Methods and Mixture Models #Machine Learning in Healthcare #cs.LG #cs.MS #stat.CO #stat.ML

paper · pdf · doi:10.48550/arxiv.2402.10797

arxiv published 2024/02/16 · arxiv updated 2024/02/22

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.

Cited by

• The Universal Warmup Path: Automatic Preconditioner Selection for HMC
• Guiding Posterior Exploration with Optimizer-Derived Geometry
• Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms
• CLAX: Fast and Flexible Neural Click Models in JAX
• Reducing normalizing flow complexity for MCMC preconditioning
• Unbiased Primordial Gravitational Wave Inference from the CMB with SMICA
• Preconditioning Hamiltonian Monte Carlo by minimizing Fisher Divergence
• Inferring Stellar Densities with Flexible Models I: The Distribution of RR Lyrae in the Milky Way with Gaia DR3
• The Sequential Monte Carlo goes NUTS: Boosting Gravitational-Wave Inference
• Jaxley: differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics
• Fisher meets Feynman: score-based variational inference with a product of experts
• Expectation-propagation for Bayesian empirical likelihood inference
• The Lensing Counter Narrative: An Effective Description of Small-Scale Clustering in Weak Lensing Power Spectra
• Argus: JAX state-space filtering for gravitational wave detection with a pulsar timing array
• Uncertainty in Physics and AI: Taxonomy, Quantification, and Validation
• High-Dimensional Bayesian Model Comparison in Cosmology with GPU-accelerated Nested Sampling and Neural Emulators
• Dynamic or Systematic? Bayesian model selection between dark energy and supernova biases
• Gravitational-wave inference at GPU speed: A bilby-like nested sampling kernel within blackjax-ns
• ParamRF: A JAX-native Framework for Declarative Circuit Modelling
• Parallelizing MCMC Across the Sequence Length
• Kilo-scale point-source inference using Parametric Cataloging
• PyBird-JAX: Accelerated inference in large-scale structure with model-independent emulation of one-loop galaxy power spectra

Discussions

• BlackJAX: Composable Bayesian Inference in Jax [hn, 3 points, 0 comments]

Related

• JAXNS: a high-performance nested sampling package based on JAX
• BAT.jl -- A Julia-based tool for Bayesian inference
• SynJax: Structured Probability Distributions for JAX
• Blang: Bayesian declarative modelling of general data structures and inference via algorithms based on distribution continua
• Past, Present and Future of Software for Bayesian Inference