The lightweight probabilistic programming library NumPyro provides a NumPy backend for Pyro (ascl:2110.016). It relies on JAX for automatic differentiation and JIT compilation to GPU/CPU. The code focuses on providing a flexible substrate for users to build on, including Pyro Primitives, inference algorithms with a particular focus on MCMC algorithms such as Hamiltonian Monte Carlo, and distribution classes, constraints and bijective transforms. NumPyro also provides effect-handlers that can be extended to implement custom inference algorithms and inference utilities.

Pyro is a flexible, scalable deep probabilistic programming library built on PyTorch. It can represent any computable probability distribution and scales to large data sets with little overhead compared to hand-written code. The library is implemented with a small core of powerful, composable abstractions. Its high-level abstractions express generative and inference models, but also allows experts to customize inference.