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 …