Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

Du Phan, Neeraj Pradhan, Martin Jankowiak
12/24/2019
stat.MLcs.AIcs.LG

Official code: pyro-ppl/numpyro (2,770 stars, Apache-2.0, JAX, last commit 2026-10-20).

Abstract

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.

View on arXivView PDF

Code Implementations(1)

pyro-ppl/numpyroOfficial100%
JAXApache-2.0

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

2,770333CSS, PythonFeb 13, 2019yesterdayApache-2.0
bayesian-inferencehmcinference-algorithmsjaxmcmc+3 more

Similar Papers

DynamicPPL: Stan-like Speed for Dynamic Probabilistic Models

87% match1 repo2892020

Efficient Probabilistic Inference in the Quest for Physics Beyond the Standard Model

81% match1 repo4022018

Inference Plans for Hybrid Particle Filtering

79% match2024

Enhancing SMT-based Weighted Model Integration by Structure Awareness

78% match1 repo262023

Scenic: A Language for Scenario Specification and Scene Generation

78% match1 repo3862018

Semirings for Probabilistic and Neuro-Symbolic Logic Programming

78% match2024

Cite this paper

@article{phan2019composable,
  title  = {Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro},
  author = {Du Phan and Neeraj Pradhan and Martin Jankowiak},
  year   = {2019},
  eprint = {1912.11554},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/1912.11554}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/1912.11554.svg)](https://codeofpaper.com/paper/1912.11554)

Discussion