mlsanityData to 7 Oct 2026

Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

arXiv 1912.11554 ↗First posted 24 Dec 2019stat.ML cs.AI cs.LGcited by 187

Not on the map. Why?

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.

Key facts. Cited by 187 papers in the map (data through 7 Oct 2026). First posted on arXiv 24 Dec 2019.

Builds on

No reference list yet.

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2021: 13; 2022: 27; 2023: 36; 2024: 41; 2025: 32; 2026: 38.

187 papers.

Sort:

By area:

  1. 2023Towards Understanding Sycophancy in Language Modelscited by 988
  2. 2021High-Dimensional Bayesian Optimization with Sparse Axis-Aligned Subspacescited by 139
  3. 2022Language Model Cascadescited by 104
  4. 2022Compositional Score Modeling for Simulation-based Inferencecited by 46
  5. 2022Robust Neural Posterior Estimation and Statistical Model Criticismcited by 45
  6. 2022Tensor Program Optimization with Probabilistic Programscited by 30
  7. 2024BlackJAX: Composable Bayesian inference in JAXcited by 24
  8. 2024sbi reloaded: a toolkit for simulation-based inference workflowscited by 23
  9. 2021A Probabilistic State Space Model for Joint Inference from Differential Equations and Datacited by 21
  10. 2022Improvement-Focused Causal Recourse (ICR)cited by 20
  11. 2022UncertaINR: Uncertainty Quantification of End-to-End Implicit Neural Representations for Computed Tomographycited by 19
  12. 2024Dissecting Human and LLM Preferencescited by 17
38 more
  1. 2024Diffusion posterior sampling for simulation-based inference in tall data settingscited by 17
  2. 2022Bayesian Active Learning for Scanning Probe Microscopy: from Gaussian Processes to Hypothesis Learningcited by 16
  3. 2021Surrogate Likelihoods for Variational Annealed Importance Samplingcited by 14
  4. 2023Fortuna: A Library for Uncertainty Quantification in Deep Learningcited by 14
  5. 2025HiBayES: A Hierarchical Bayesian Modeling Framework for AI Evaluation Statisticscited by 14
  6. 2021Gaussian processes meet NeuralODEs: A Bayesian framework for learning the dynamics of partially observed systems from scarce and noisy datacited by 13
  7. 2023Detect Everything with Few Examplescited by 13
  8. 2024Scalable Bayesian Learning with posteriorscited by 13
  9. 2026Diagnosing the Reliability of LLM-as-a-Judge via Item Response Theorycited by 13
  10. 2026EdgeBench: Unveiling Scaling Laws of Learning from Real-World Environmentscited by 12
  11. 2021Efficient and Generalizable Tuning Strategies for Stochastic Gradient MCMCcited by 11
  12. 2021PriorVAE: Encoding spatial priors with VAEs for small-area estimationcited by 11
  13. 2023Choosing a Proxy Metric from Past Experimentscited by 11
  14. 2021ProbNum: Probabilistic Numerics in Pythoncited by 10
  15. 2022Deterministic Langevin Monte Carlo with Normalizing Flows for Bayesian Inferencecited by 10
  16. 2024Predicting Emergent Capabilities by Finetuningcited by 10
  17. 2023Physics-informed Information Field Theory for Modeling Physical Systems with Uncertainty Quantificationcited by 9
  18. 2024GIST: Gibbs self-tuning for locally adaptive Hamiltonian Monte Carlocited by 9
  19. 2024Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformerscited by 9
  20. 2023Misspecification-robust Sequential Neural Likelihood for Simulation-based Inferencecited by 8
  21. 2023Bayesian neural networks via MCMC: a Python-based tutorialcited by 8
  22. 2024Automated Discovery of Pairwise Interactions from Unstructured Datacited by 8
  23. 2022Does the Market of Citations Reward Reproducible Work?cited by 7
  24. 2023On the Properties and Estimation of Pointwise Mutual Information Profilescited by 7
  25. 2023SigFormer: Signature Transformers for Deep Hedgingcited by 7
  26. 2024Demonstration of Robust and Efficient Quantum Property Learning with Shallow Shadowscited by 7
  27. 2022PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAXcited by 6
  28. 2022Strong Lensing Source Reconstruction Using Continuous Neural Fieldscited by 6
  29. 2022An Ordinal Latent Variable Model of Conflict Intensitycited by 6
  30. 2023Fluctuation without dissipation: Microcanonical Langevin Monte Carlocited by 6
  31. 2023PriorCVAE: scalable MCMC parameter inference with Bayesian deep generative modellingcited by 6
  32. 2023Uncertainty and Structure in Neural Ordinary Differential Equationscited by 6
  33. 2022Loss-calibrated expectation propagation for approximate Bayesian decision-makingcited by 5
  34. 2022Active learning in open experimental environments: selecting the right information channel(s) based on predictability in deep kernel learningcited by 5
  35. 2022Neural Implicit Manifold Learning for Topology-Aware Density Estimationcited by 5
  36. 2023PINN surrogate of Li-ion battery models for parameter inference. Part II: Regularization and application of the pseudo-2D modelcited by 5
  37. 2024Stable Training of Normalizing Flows for High-dimensional Variational Inferencecited by 5
  38. 2024Simulation-based Inference with the Python Package sbijaxcited by 5

50 most cited of 187 All 187 (JSON)