BlackJAX: Library of samplers for JAX

BlackJAX is a sampling library designed for ease of use, speed, and modularity and works on CPU as well as GPU. It is not a probabilistic programming library (PLL), though it integrates well with PPLs as long as they can provide a (potentially unnormalized) log-probability density function compatible with JAX. BlackJAX is written in pure Python and depends on XLA via JAX (ascl:2111.002). It can be used by those who have a logpdf and need a sampler or need more than a general-purpose sampler. It is also useful for building a sample on GPU and for users who want to learn how sampling algorithms work.