mlsanityData to 14 Oct 2026

BlackJAX: Composable Bayesian inference in JAX

arXiv 2402.10797 ↗First posted 16 Feb 2024cs.MS cs.LG stat.COcited by 26

Not on the map. Why?

Abstract (arXiv)

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.

Key facts. Cited by 26 papers in the map (data through 14 Oct 2026). Builds on 9 of its 45 references that are in the map. First posted on arXiv 16 Feb 2024.

Builds on

9 of 45 references.

References · 9

  1. 2018Pyro: Deep Universal Probabilistic Programmingcited by 380
  2. 2019Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyrocited by 191
  3. 2018Bayesian Layers: A Module for Neural Network Uncertaintycited by 57
  4. 2021Bayes-Newton Methods for Approximate Bayesian Inference with PSD Guaranteescited by 13
  5. 2022Transport Elliptical Slice Samplingcited by 12
  6. 2016Identifying the Optimal Integration Time in Hamiltonian Monte Carlocited by 4
  7. 2020Non-reversibly updating a uniform [0,1] value for Metropolis accept/reject decisionscited by 3
  8. 2022Benchmarking Bayesian neural networks and evaluation metrics for regression taskscited by 3
  9. 2023Bayesian cross-validation by parallel Markov Chain Monte Carlocited by 3

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 6; 2025: 6; 2026: 14.

26 papers.

Sort:

By area:

  1. 2024Sequential Controlled Langevin Diffusionscited by 48
  2. 2026BayesFlow 2: Multi-Backend Amortized Bayesian Inference in Pythoncited by 9
  3. 2024Simulation-based Inference with the Python Package sbijaxcited by 6
  4. 2026On the Interplay of Priors and Overparametrization in Bayesian Neural Network Posteriorscited by 3
  5. 2026Uncertainty in Physics and AI: Taxonomy, Quantification, and Validationcited by 3
  6. 2025Geodesic Slice Sampler for Multimodal Distributions with Strong Curvaturecited by 2
  7. 2025Fisher meets Feynman: score-based variational inference with a product of expertscited by 2
  8. 2026Statistical Learning Analysis of Physics-Informed Neural Networkscited by 1
  9. 2026How to Approximate Inference with Subtractive Mixture Modelscited by 1
  10. 2026bde: A Python Package for Bayesian Deep Ensembles via MILEcited by 1
  11. 2024Federated Learning for Non-factorizable Models using Deep Generative Prior Approximationscited by 0
  12. 2024Diffusion Generative Modelling for Divide-and-Conquer MCMCcited by 0
14 more
  1. 2024Scalable Monte Carlo for Bayesian Learningcited by 0
  2. 2024Flow Annealed Importance Sampling Bootstrap meets Differentiable Particle Physicscited by 0
  3. 2025Efficiently Vectorized MCMC on Modern Acceleratorscited by 0
  4. 2025Bayesian Inverse Physics for Neuro-Symbolic Robot Learningcited by 0
  5. 2025Reducing normalizing flow complexity for MCMC preconditioningcited by 0
  6. 2025CLAX: Fast and Flexible Neural Click Models in JAXcited by 0
  7. 2026Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamicscited by 0
  8. 2026Corrected Integrated Laplace Approximation for Bayesian Inference in Latent Gaussian Modelscited by 0
  9. 2026GenSBI: Generative Methods for Simulation-Based Inference in JAXcited by 0
  10. 2026Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveformscited by 0
  11. 2026Guiding Posterior Exploration with Optimizer-Derived Geometrycited by 0
  12. 2026Subspace Inference Enables Efficient Active Reward Learning from Preferencescited by 0
  13. 2026Multiple latent orderings better predict language model preferencescited by 0
  14. 2026Predictively Oriented Gaussian Process Posteriorscited by 0