Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

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.

10 pages, 2 figures; NeurIPS 2019 Program Transformations for Machine Learning Workshop

References in corpus (2)

• Effect Handling for Composable Program Transformations in Edward2
• The Kernel Interaction Trick: Fast Bayesian Discovery of Pairwise Interactions in High Dimensions

Cited by in corpus (62)

• Hypothesis Learning in Automated Experiment: Application to Combinatorial Materials Libraries
• JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
• High-Dimensional Bayesian Optimization with Sparse Axis-Aligned Subspaces
• CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators
• Bayesian neural networks via MCMC: a Python-based tutorial
• Autonomous scanning probe microscopy with hypothesis learning: Exploring the physics of domain switching in ferroelectric materials
• Auto-Differentiable Spectrum Model for High-Dispersion Characterization of Exoplanets and Brown Dwarfs
• The evolution of the Milky Way's thin disc radial metallicity gradient with K2 asteroseismic ages
• Three-dimensional dust density structure of the Orion, Cygnus X,Taurus, and Perseus star-forming regions
• Fast optical flares from M dwarfs detected by a one-second-cadence survey with Tomo-e Gozen
• Rethinking Parameter Counting in Deep Models: Effective Dimensionality Revisited
• Re-Envisioning Numerical Information Field Theory (NIFTy.re): A Library for Gaussian Processes and Variational Inference
• candl: Cosmic Microwave Background Analysis with a Differentiable Likelihood
• Modelling Cosmic Radiation Events in the Tree-ring Radiocarbon Record
• The dawn of disks: unveiling the turbulent ionised gas kinematics of the galaxy population at $z\sim 4-6$ with JWST/NIRCam grism spectroscopy
• On the functional form of the radial acceleration relation
• Conditional independence by typing
• Atacama Cosmology Telescope measurements of a large sample of candidates from the Massive and Distant Clusters of WISE Survey: Sunyaev-Zeldovich effect confirmation of MaDCoWS candidates using ACT
• Structured Variational Inference for Simulating Populations of Radio Galaxies
• Marginalised Normal Regression: Unbiased curve fitting in the presence of x-errors
• Estimating Model Uncertainty of Neural Networks in Sparse Information Form
• Precision Ephemerides for Gravitational-wave Searches -- IV: Corrected and refined ephemeris for Scorpius X-1
• DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology
• Two per cent measurement of $H_0$ from Cepheids alone
• On the Evolution of Rotational Modulation Amplitude in Solar-mass Main-sequence Stars
• Joint Distributions for TensorFlow Probability
• Automatic structured variational inference
• Cosmic Cartography: Bayesian reconstruction of the galaxy density informed by large-scale structure
• SIDE-real: Supernova Ia Dust Extinction with truncated marginal neural ratio estimation applied to real data
• Differentiable Optics with dLux I: Deep calibration of Flat Field and Phase Retrieval with Automatic Differentiation
• SigFormer: Signature Transformers for Deep Hedging
• Disentangling transients and their host galaxies with Scarlet2: A framework to forward model multi-epoch imaging
• Map-level baryonification: unified treatment of weak lensing two-point and higher-order statistics
• DAmodel: Hierarchical Bayesian Modelling of DA White Dwarfs for Spectrophotometric Calibration
• Clustering properties of the CatWISE2020 quasar catalogue and their impact on the cosmic dipole anomaly
• A Constrained NILC method for CMB B mode observations
• Testing the local supervoid solution to the Hubble tension with direct distance tracers
• The Velocity Field Olympics: Assessing velocity field reconstructions with direct distance tracers
• Mixed Hamiltonian Monte Carlo for Mixed Discrete and Continuous Variables
• The subtle statistics of the distance ladder: On the distance prior and selection effects
• No evidence for local $H_0$ anisotropy from Tully--Fisher or supernova distances
• DeepSSM: an emulator of gravitational wave spectra from sound waves during cosmological first-order phase transitions
• Hierarchical Bayesian estimation of motor-evoked potential recruitment curves yields accurate and robust estimates
• A Probabilistic State Space Model for Joint Inference from Differential Equations and Data
• Constraining the primordial power spectrum using a differentiable likelihood
• Characterising the Standardisation Properties of Type Ia Supernovae in the z band with Hierarchical Bayesian Modelling
• On decision-theoretic model assessment for structural deterioration monitoring
• Forward-modelling Milky Way Cepheids: selection effects and physical priors in the Gaia-HST calibration
• Measuring the Stellar-to-Halo Mass Relation at $\sim 10^{10}$ Solar masses, using forthcoming space-based imaging of galaxy-galaxy strong lenses
• Hamiltonian Monte Carlo for Regression with High-Dimensional Categorical Data
• Effects of Antivaccine Tweets on COVID-19 Vaccinations, Cases, and Deaths
• D3p -- A Python Package for Differentially-Private Probabilistic Programming
• Locally Differentially Private Bayesian Inference
• Dippers from TESS Full-frame Images. II. Spectroscopic Characterization of Four Young Dippers
• Automatic Guide Generation for Stan via NumPyro
• FunMC: A functional API for building Markov Chains
• The GALAH Survey: Neutron-Capture Elemental Abundances for 350,000 Gaia-RVS spectra and the Chemodynamics of Accreted Structures
• First constraints on the local ionization topology in front of two quasars at z ~ 7.5
• $S_8$ from peculiar velocities: agreement with Planck for Tully--Fisher and supernovae, tension for the fundamental plane
• FlowSN: Neural Simulation-Based Inference under Realistic Selection Effects applied to Supernova Cosmology
• Neural network emulator to constrain the high-$z$ IGM thermal state from Lyman-$\alpha$ forest flux auto-correlation function
• Automatically Differentiable Random Coefficient Logistic Demand Estimation