BlackJAX: Composable Bayesian inference in JAX

arXiv:2402.10797v1 Announce Type: cross
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 …

abstract algorithms arxiv bayesian bayesian inference computation cs.lg cs.ms functional implementation inference jax library python sampling speed stat.co stat.ml type