# Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

arXiv: https://arxiv.org/abs/1912.11554 · HTML: https://mlsanity.com/paper/1912.11554/

## Key facts

- Cited by 191 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 24 Dec 2019.

## Built on by (most cited first, 30 of 191)

- [Towards Understanding Sycophancy in Language Models](https://mlsanity.com/paper/2310.13548/) (2023)
- [High-Dimensional Bayesian Optimization with Sparse Axis-Aligned Subspaces](https://mlsanity.com/paper/2103.00349/) (2021)
- [Language Model Cascades](https://mlsanity.com/paper/2207.10342/) (2022)
- [Compositional Score Modeling for Simulation-based Inference](https://mlsanity.com/paper/2209.14249/) (2022)
- [Robust Neural Posterior Estimation and Statistical Model Criticism](https://mlsanity.com/paper/2210.06564/) (2022)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [BlackJAX: Composable Bayesian inference in JAX](https://mlsanity.com/paper/2402.10797/) (2024)
- [sbi reloaded: a toolkit for simulation-based inference workflows](https://mlsanity.com/paper/2411.17337/) (2024)
- [A Probabilistic State Space Model for Joint Inference from Differential Equations and Data](https://mlsanity.com/paper/2103.10153/) (2021)
- [Improvement-Focused Causal Recourse (ICR)](https://mlsanity.com/paper/2210.15709/) (2022)
- [EdgeBench: Unveiling Scaling Laws of Learning from Real-World Environments](https://mlsanity.com/paper/2607.05155/) (2026)
- [UncertaINR: Uncertainty Quantification of End-to-End Implicit Neural Representations for Computed Tomography](https://mlsanity.com/paper/2202.10847/) (2022)
- [Diffusion posterior sampling for simulation-based inference in tall data settings](https://mlsanity.com/paper/2404.07593/) (2024)
- [HiBayES: A Hierarchical Bayesian Modeling Framework for AI Evaluation Statistics](https://mlsanity.com/paper/2505.05602/) (2025)
- [Dissecting Human and LLM Preferences](https://mlsanity.com/paper/2402.11296/) (2024)
- [Bayesian Active Learning for Scanning Probe Microscopy: from Gaussian Processes to Hypothesis Learning](https://mlsanity.com/paper/2205.15458/) (2022)
- [Diagnosing the Reliability of LLM-as-a-Judge via Item Response Theory](https://mlsanity.com/paper/2602.00521/) (2026)
- [Surrogate Likelihoods for Variational Annealed Importance Sampling](https://mlsanity.com/paper/2112.12194/) (2021)
- [Fortuna: A Library for Uncertainty Quantification in Deep Learning](https://mlsanity.com/paper/2302.04019/) (2023)
- [Detect Everything with Few Examples](https://mlsanity.com/paper/2309.12969/) (2023)
- [Gaussian processes meet NeuralODEs: A Bayesian framework for learning the dynamics of partially observed systems from scarce and noisy data](https://mlsanity.com/paper/2103.03385/) (2021)
- [Scalable Bayesian Learning with posteriors](https://mlsanity.com/paper/2406.00104/) (2024)
- [Efficient and Generalizable Tuning Strategies for Stochastic Gradient MCMC](https://mlsanity.com/paper/2105.13059/) (2021)
- [PriorVAE: Encoding spatial priors with VAEs for small-area estimation](https://mlsanity.com/paper/2110.10422/) (2021)
- [Choosing a Proxy Metric from Past Experiments](https://mlsanity.com/paper/2309.07893/) (2023)
- [ProbNum: Probabilistic Numerics in Python](https://mlsanity.com/paper/2112.02100/) (2021)
- [Deterministic Langevin Monte Carlo with Normalizing Flows for Bayesian Inference](https://mlsanity.com/paper/2205.14240/) (2022)
- [Predicting Emergent Capabilities by Finetuning](https://mlsanity.com/paper/2411.16035/) (2024)
- [Physics-informed Information Field Theory for Modeling Physical Systems with Uncertainty Quantification](https://mlsanity.com/paper/2301.07609/) (2023)
- [GIST: Gibbs self-tuning for locally adaptive Hamiltonian Monte Carlo](https://mlsanity.com/paper/2404.15253/) (2024)

## Abstract (arXiv)

NumPyro is a lightweight library that provides an alternate NumPy backend to the Pyro probabilistic programming language with the same modeling interface, language primitives and effect handling abstractions. Effect handlers allow Pyro's modeling API to be extended to NumPyro despite its being built atop a fundamentally different JAX-based functional backend. In this work, we demonstrate the power of composing Pyro's effect handlers with the program transformations that enable hardware acceleration, automatic differentiation, and vectorization in JAX. In particular, NumPyro provides an iterative formulation of the No-U-Turn Sampler (NUTS) that can be end-to-end JIT compiled, yielding an implementation that is much faster than existing alternatives in both the small and large dataset regimes.