Pyro: Deep Universal Probabilistic Programming for PyTorch

Jul 10, 2025

Introduction

Modern machine learning often struggles with uncertainty. While standard deep learning models provide point estimates, they frequently fail to quantify how sure they are about a prediction. Pyro solves this by unifying the best of modern deep learning with Bayesian modeling, allowing developers to build models that reason under uncertainty. Built on PyTorch, Pyro provides a scalable framework for deep probabilistic programming, enabling the creation of complex stochastic models that can be trained on massive datasets.

What Is Pyro?

Pyro is a deep universal probabilistic programming language (PPL) written in Python and supported by PyTorch on the backend. It is designed for researchers and developers who need to define complex probabilistic models using Python code and perform automated Bayesian inference to reason under uncertainty.

Originally developed at Uber AI and now a project of the Linux Foundation (LF AI & Data), Pyro is licensed under the Apache 2.0 License. It allows users to combine the flexibility of Python with the tensor computations and automatic differentiation of PyTorch, making it a powerful tool for both statistical inference and deep generative modeling.

Why Pyro Matters

Traditional Bayesian methods often struggle to scale to large datasets or integrate with deep neural networks. Pyro fills this gap by leveraging PyTorch’s GPU acceleration and scalable optimizers. This allows practitioners to move beyond simple conjugate priors and implement Stochastic Variational Inference (SVI), which can handle millions of data points with minimal overhead.

The project has gained significant traction in the AI community, being adopted by institutions like the Broad Institute and Uber. By providing a universal PPL, Pyro ensures that any computable probability distribution can be represented, removing the limitations found in more rigid probabilistic languages.

For developers already familiar with the PyTorch ecosystem, Pyro is the natural choice for adding probabilistic layers to their existing deep learning pipelines, enabling a transition from deterministic models to those that quantify uncertainty.

Key Features

  • Universal Probabilistic Programming: Pyro can represent any computable probability distribution, allowing for the creation of arbitrary stochastic processes without being limited by a specific domain-specific language.
  • Scalable Inference Engines: The library provides automated inference algorithms, including Stochastic Variational Inference (SVI) and Markov Chain Monte Carlo (MCMC), which scale to large datasets using PyTorch’s backend.
  • Deep Learning Integration: Because it is built on PyTorch, Pyro allows for the seamless combination of deep neural networks with probabilistic models, enabling the creation of Deep Generative Models like Variational Autoencoders (VAEs).
  • Flexible Modeling Primitives: Using pyro.sample, developers can define latent random variables, while pyro.plate enables efficient vectorization and declares independence across data points.
  • Automated Exact Inference: Pyro provides specialized tools for discrete latent variable models, such as Hidden Markov Models (HMMs), allowing for automated exact inference.
  • Composable Abstractions: The core is implemented with a small set of powerful, composable abstractions that keep the library agile and maintainable while providing high-level control for experts.

How Pyro Compares

Pyro competes with other probabilistic programming languages like PyMC and TensorFlow Probability. While all three allow for Bayesian inference, their backends and target audiences differ significantly.

Feature Pyro PyMC TF Probability
Backend PyTorch PyTensor / Aesara TensorFlow
Primary Inference SVI / MCMC MCMC (NUTS) SVI / MCMC
Deep Learning Integration Native Limited Native
Scalability High (GPU) Moderate High (GPU)

Pyro’s primary differentiator is its deep integration with PyTorch. For users who are already invested in the PyTorch ecosystem, Pyro provides a much smoother experience than PyMC, which relies on its own tensor library. While TensorFlow Probability is a similar competitor, Pyro’s “universal” nature—allowing arbitrary Python code within models—makes it more flexible for researchers who need to customize their generative processes.

A notable tradeoff is that PyMC is often seen as more intuitive for traditional statisticians who prefer a more declarative style of modeling. Pyro requires a more programmatic approach, which is more powerful but has a steeper learning curve for those not familiar with stochastic functions.

Getting Started: Installation

Pyro requires PyTorch to be installed first. Ensure you have a compatible version of PyTorch installed on your system before proceeding with the Pyro installation.

Install via Pip

The fastest way to install the stable release of Pyro is through pip:

pip install pyro-ppl

Install from Source

If you need the latest features from the master branch, you can install from source:

git clone https://github.com/pyro-ppl/pyro.git
cd pyro
pip install .

Install with Extras

To install the dependencies required to run the probabilistic models included in the examples and tutorials, use the following command:

pip install pyro-ppl[extras]

How to Use Pyro

In Pyro, models are written as stochastic functions. A model is a Python function that uses Pyro primitives to define a generative process. The most basic workflow involves defining a model and a “guide” (the variational distribution used for inference).

To start, you use pyro.sample to define random variables. You can name these variables, which allows Pyro to track them across the model and the guide. By using pyro.plate, you can tell Pyro that certain observations are independent and identically distributed (i.i.d.), which allows the backend to vectorize the computations for performance.

Once the model and guide are defined, you use an inference engine like SVI (Stochastic Variational Inference) to optimize the guide parameters so that the variational distribution approximates the posterior distribution of the latent variables.

Code Examples

The following examples demonstrate how to define a simple probabilistic model in Pyro. These examples are based on the official repository’s introductory materials.

Simple Bayesian Coin Toss

This example shows how to define a model for a coin toss where the probability of heads is a latent variable.

import torch
import pyro
import pyro.distributions as dist

def model(data):
    # Prior for the probability of heads
    heads_prob = pyro.sample("heads_prob", dist.Beta(1.0, 1.0))
    
    # Likelihood of the observed data
    with pyro.plate("data", len(data)):
        pyro.sample("obs", dist.Bernoulli(heads_prob), obs=data)

# Example data: 10 heads, 2 tails
data = torch.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0])
model(data)

In this snippet, pyro.sample creates a latent random variable heads_prob with a Beta distribution. The pyro.plate block ensures that the coin tosses are treated as independent observations.

Stochastic Variational Inference (SVI)

SVI is the primary way to perform inference in Pyro. Here is how you set up a basic SVI loop.

from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam

# Define a guide (the variational distribution)
def guide(data):
    # The guide must use the same names for random variables as the model
    heads_prob_loc = pyro.param("heads_prob_loc", torch.tensor(0.0))
    heads_prob_scale = pyro.param("heads_prob_scale", torch.tensor(1.0))
    
    # Sample from the variational distribution
    pyro.sample("heads_prob", dist.Normal(heads_prob_loc, heads_prob_scale))

# Setup SVI
svi = SVI(model, guide, Adam(), loss=Trace_ELBO())

# Optimization loop
for step in range(1000):
    svi.step(data)

This example demonstrates the use of pyro.param to register learnable parameters that the SVI engine will optimize to minimize the Evidence Lower Bound (ELBO).

Real-World Use Cases

Pyro is particularly effective in scenarios where uncertainty quantification is critical and datasets are large.

  • Time Series Forecasting: Using Hierarchical Models, Pyro can be used to forecast demand for thousands of products across different regions, allowing the model to share information between similar products (hierarchical priors).
  • Deep Generative Models: Researchers use Pyro to build Variational Autoencoders (VAEs) and other deep generative models that combine neural networks with probabilistic latent spaces to generate new data or perform dimensionality reduction.
  • Experimental Design: In scientific modeling, Pyro can be used to optimize the same process of data collection, allowing researchers to determine which experiments to run next to maximize the information gain.
  • Epidemiology: By modeling the spread of a disease using stochastic differential equations or discrete-time Markov models, Pyro can be used to estimate parameters of infection rates and recovery rates under uncertainty.

Contributing to Pyro

Pyro is an open-source project under the Linux Foundation and welcomes contributions from the community. Developers can contribute by improving documentation, fixing bugs, or implementing new distributions and inference algorithms.

To contribute, start by opening a GitHub issue to discuss proposed changes, especially for large design changes. For small fixes, developers can submit a pull request directly. The project maintains a Technical Steering Committee (TSC) that governs development and core architectural decisions.

Contributions should follow the project’s coding standards and ensure that tests are run locally before submission. The project encourages the use of GitHub Discussions for community interaction and the laziest way to contribute is to start by tackling “good first issues” labeled in the repository.

Community and Support

Pyro has a dedicated community of researchers and researchers. The primary hub for support is the official Pyro forum, where users can ask questions and discuss modeling challenges. The laziest way to get started is to explore the official tutorials, which are comprehensive and build up from basic Bayesian regression to complex deep generative models.

Official channels include the GitHub repository for issue tracking and the Linux Foundation’s LF AI & Data project page. The community is active in sharing tutorials and community-led notebooks, such as those translating the “Statistical Rethinking” book into Pyro code.

Conclusion

Pyro is the premier choice for developers who need to combine the power of PyTorch with the rigor of Bayesian inference. It is a universal probabilistic programming language that scales to large datasets and integrates seamlessly with deep learning pipelines. While it has a steeper learning curve than some traditional Bayesian tools, the flexibility and scalability it provides is unmatched for modern AI applications.

If you are building a system that needs to reason under uncertainty, quantify risk, or build deep generative models, Pyro is the right tool. We recommend starting with the official tutorials and then experimenting with the SVI engine to move your deterministic models into the probabilistic realm.

Star the repo, try the quickstart, and join the community forum to start building probabilistic AI.

What is Pyro and what problem does it solve?

Pyro is a deep universal probabilistic programming language built on PyTorch that allows developers to create complex probabilistic models and perform Bayesian inference. It solves the problem of uncertainty quantification in machine learning, allowing models to provide a distribution of possible outcomes rather than a single point estimate.

How do I install Pyro?

Pyro can be installed via pip using the command pip install pyro-ppl. It requires PyTorch to be installed first on your system. For users who need additional dependencies for tutorials and examples, pip install pyro-ppl[extras] is recommended.

How does Pyro compare to PyMC?

Pyro is built on PyTorch, making it more scalable for large datasets and better integrated with deep learning. PyMC is a more traditional Bayesian tool that uses its own tensor library (PyTensor), and while it is more intuitive for some statisticians, it is not as natively integrated with deep learning frameworks.

Can I use Pyro for deep generative models like VAEs?

Pyro is specifically designed for this. By combining pyro.sample and SVI, developers can build Variational Autoencoders (VAEs) and other deep generative models that combine neural networks with probabilistic latent spaces.

What is the difference between Pyro and NumPyro?

Pyro is backed by PyTorch, while NumPyro is a NumPy-backed version of Pyro that uses JAX for automatic differentiation and JIT compilation. NumPyro is often significantly faster for MCMC inference but may have different integration paths for deep learning pipelines.

Is Pyro open source?

Pyro is an Apache 2.0-Licensed open source project and is now a project of the Linux Foundation (LF AI & Data).

What are the main inference engines in Pyro?

The primary inference engines are Stochastic Variational Inference (SVI) and Markov Chain Monte Carlo (MCMC), which allow for the approximate estimation of the posterior distributions of latent variables.