0
# Pyro
1

2
Pyro is a flexible, scalable deep probabilistic programming library built on PyTorch that enables universal probabilistic modeling and inference. It combines the expressiveness of probabilistic programming with the power of deep learning, providing a comprehensive toolkit for Bayesian modeling, variational inference, and uncertainty quantification.
3

4
## Package Information
5

6
- **Package Name**: pyro-ppl
7
- **Package Type**: pypi
8
- **Language**: Python
9
- **Installation**: `pip install pyro-ppl`
10
- **Requirements**: Python 3.8+, PyTorch 2.0+
11

12
## Core Imports
13

14
```python
15
import pyro
16
```
17

18
Common imports for probabilistic programming:
19

20
```python
21
import pyro
22
import pyro.distributions as dist
23
from pyro.infer import SVI, Trace_ELBO
24
from pyro.optim import Adam
25
import torch
26
```
27

28
## Basic Usage
29

30
```python
31
import pyro
32
import pyro.distributions as dist
33
from pyro.infer import SVI, Trace_ELBO
34
from pyro.optim import Adam
35
import torch
36

37
# Define a simple Bayesian model
38
def model(data):
39
    # Prior on the parameter
40
    theta = pyro.sample("theta", dist.Beta(1.0, 1.0))
41
    
42
    # Likelihood
43
    with pyro.plate("data", len(data)):
44
        pyro.sample("obs", dist.Bernoulli(theta), obs=data)
45

46
# Define a variational guide (posterior approximation)
47
def guide(data):
48
    # Variational parameters
49
    alpha_q = pyro.param("alpha_q", torch.tensor(1.0), constraint=dist.constraints.positive)
50
    beta_q = pyro.param("beta_q", torch.tensor(1.0), constraint=dist.constraints.positive)
51
    
52
    # Variational distribution
53
    pyro.sample("theta", dist.Beta(alpha_q, beta_q))
54

55
# Generate some data
56
data = torch.tensor([1.0, 0.0, 1.0, 1.0, 0.0])
57

58
# Set up stochastic variational inference
59
svi = SVI(model, guide, Adam({"lr": 0.01}), loss=Trace_ELBO())
60

61
# Training loop
62
for step in range(1000):
63
    loss = svi.step(data)
64
    if step % 100 == 0:
65
        print(f"Step {step}, Loss: {loss}")
66

67
# Get posterior samples
68
from pyro.infer import Predictive
69
predictive = Predictive(model, guide=guide, num_samples=1000)
70
samples = predictive(data)
71
print(f"Posterior mean of theta: {samples['theta'].mean():.3f}")
72
```
73

74
## Architecture
75

76
Pyro's architecture is built on several key design principles:
77

78
- **Universal**: Can represent any computable probability distribution
79
- **Scalable**: Efficient GPU computation and minibatch training via PyTorch
80
- **Minimal**: Small core with composable primitives
81
- **Flexible**: High-level abstractions with low-level customization
82

83
### Core Components
84

85
- **Probabilistic Programs**: Functions that use `pyro.sample` and `pyro.param`
86
- **Effect Handlers (Poutine)**: Composable program transformations for inference
87
- **Distributions**: 100+ probability distributions with automatic differentiation
88
- **Inference**: Variational inference, MCMC, and importance sampling algorithms
89
- **Neural Networks**: Deep probabilistic models with `pyro.nn`
90

91
## Capabilities
92

93
### Core Probabilistic Programming
94

95
Primary functions and constructs for building probabilistic programs, including sampling, parameter management, and independence declarations.
96

97
```python { .api }
98
def sample(name: str, fn: dist.Distribution, *args, obs=None, obs_mask=None, infer=None, **kwargs):
99
    """
100
    Primitive stochastic function for probabilistic programming.
101
    
102
    Parameters:
103
    - name (str): Unique name for the sample site
104
    - fn (Distribution): Probability distribution to sample from
105
    - obs (Tensor, optional): Observed data to condition on
106
    - obs_mask (Tensor, optional): Mask for observed data
107
    - infer (dict, optional): Inference configuration
108
    
109
    Returns:
110
    Tensor: Sample from the distribution
111
    """
112

113
def param(name: str, init_tensor=None, constraint=None, event_dim=None):
114
    """
115
    Declare and retrieve learnable parameters.
116
    
117
    Parameters:
118
    - name (str): Parameter name
119
    - init_tensor (Tensor, optional): Initial value
120
    - constraint (Constraint, optional): Parameter constraint
121
    - event_dim (int, optional): Event dimension
122
    
123
    Returns:
124
    Tensor: Parameter tensor
125
    """
126

127
def plate(name: str, size: int, subsample_size=None, dim=None):
128
    """
129
    Independence context manager for vectorized computation.
130
    
131
    Parameters:
132
    - name (str): Plate name
133
    - size (int): Plate size
134
    - subsample_size (int, optional): Subsample size for minibatching
135
    - dim (int, optional): Tensor dimension
136
    
137
    Returns:
138
    PlateMessenger: Context manager
139
    """
140

141
def factor(log_factor):
142
    """
143
    Add arbitrary factor to log probability.
144
    
145
    Parameters:
146
    - log_factor (Tensor): Log probability factor
147
    """
148
```
149

150
[Core Programming](./core-programming.md)
151

152
### Probability Distributions
153

154
Comprehensive collection of probability distributions including continuous, discrete, multivariate, and specialized distributions for probabilistic modeling.
155

156
```python { .api }
157
# Core continuous distributions
158
class Normal(dist.Distribution):
159
    def __init__(self, loc: torch.Tensor, scale: torch.Tensor): ...
160

161
class Beta(dist.Distribution):
162
    def __init__(self, concentration1: torch.Tensor, concentration0: torch.Tensor): ...
163

164
class Gamma(dist.Distribution):
165
    def __init__(self, concentration: torch.Tensor, rate: torch.Tensor): ...
166

167
# Discrete distributions
168
class Bernoulli(dist.Distribution):
169
    def __init__(self, probs: torch.Tensor = None, logits: torch.Tensor = None): ...
170

171
class Categorical(dist.Distribution):
172
    def __init__(self, probs: torch.Tensor = None, logits: torch.Tensor = None): ...
173

174
# Multivariate distributions
175
class MultivariateNormal(dist.Distribution):
176
    def __init__(self, loc: torch.Tensor, covariance_matrix: torch.Tensor = None, 
177
                 precision_matrix: torch.Tensor = None, scale_tril: torch.Tensor = None): ...
178
```
179

180
[Distributions](./distributions.md)
181

182
### Inference Methods
183

184
Scalable inference algorithms including variational inference (SVI), Markov Chain Monte Carlo (MCMC), and importance sampling for posterior approximation.
185

186
```python { .api }
187
class SVI:
188
    """Stochastic Variational Inference."""
189
    def __init__(self, model, guide, optim, loss):
190
        """
191
        Parameters:
192
        - model: Generative model function
193
        - guide: Variational guide function  
194
        - optim: Optimizer instance
195
        - loss: Loss function (typically ELBO)
196
        """
197
    
198
    def step(self, *args, **kwargs) -> float:
199
        """Perform one SVI step, returns loss."""
200

201
class MCMC:
202
    """Markov Chain Monte Carlo."""
203
    def __init__(self, kernel, num_samples: int, warmup_steps: int = None): ...
204
    
205
    def run(self, *args, **kwargs): ...
206

207
class Predictive:
208
    """Generate predictions from posterior samples."""
209
    def __init__(self, model, guide=None, posterior_samples=None, num_samples=None): ...
210
```
211

212
[Inference](./inference.md)
213

214
### Neural Networks Integration  
215

216
Deep probabilistic models combining neural networks with probabilistic programming, including Bayesian neural networks and stochastic layers.
217

218
```python { .api }
219
class PyroModule(torch.nn.Module):
220
    """Base class for Pyro modules with parameter/sample integration."""
221
    
222
class PyroParam:
223
    """Descriptor for Pyro parameters in modules."""
224
    def __init__(self, init_tensor, constraint=None, event_dim=None): ...
225

226
class PyroSample:
227
    """Descriptor for Pyro samples in modules."""  
228
    def __init__(self, prior): ...
229

230
class DenseNN(PyroModule):
231
    """Dense neural network for use in flows and guides."""
232
    def __init__(self, input_dim: int, hidden_dims: List[int], output_dim: int): ...
233
```
234

235
[Neural Networks](./neural-networks.md)
236

237
### Transforms and Constraints
238

239
Bijective transformations and parameter constraints for reparametrization and constrained optimization in probabilistic models.
240

241
```python { .api }
242
# Core transforms
243
class Transform:
244
    """Base class for bijective transforms."""
245
    def __call__(self, x): ...
246
    def inv(self, y): ...
247
    def log_abs_det_jacobian(self, x, y): ...
248

249
class AffineTransform(Transform):
250
    def __init__(self, loc: torch.Tensor, scale: torch.Tensor): ...
251

252
# Flow-based transforms  
253
class AffineAutoregressive(Transform):
254
    def __init__(self, autoregressive_nn, log_scale_min_clip: float = -5.0): ...
255

256
# Constraints
257
class Constraint:
258
    """Base class for parameter constraints."""
259
    def check(self, value): ...
260

261
positive: Constraint
262
unit_interval: Constraint  
263
simplex: Constraint
264
```
265

266
[Transforms and Constraints](./transforms-constraints.md)
267

268
### Gaussian Processes
269

270
Gaussian process models for non-parametric Bayesian modeling, including kernels, likelihoods, and efficient GP inference.
271

272
```python { .api }
273
class Kernel:
274
    """Base class for GP kernels."""
275
    def forward(self, X, Z=None, diag: bool = False): ...
276

277
class RBF(Kernel):
278
    """Radial Basis Function kernel."""
279
    def __init__(self, input_dim: int, lengthscale=None, variance=None): ...
280

281
class GPModel:
282
    """Base Gaussian Process model."""
283
    def __init__(self, X, y, kernel, likelihood): ...
284
```
285

286
[Gaussian Processes](./gaussian-processes.md)
287

288
### Optimization
289

290
Optimization utilities and PyTorch optimizer wrappers for training probabilistic models with Pyro's parameter store system.
291

292
```python { .api }
293
class PyroOptim:
294
    """Base wrapper for PyTorch optimizers."""
295
    def __init__(self, optim_constructor, optim_args, clip_args=None): ...
296

297
class ClippedAdam(PyroOptim):
298
    """Adam optimizer with gradient clipping."""
299
    def __init__(self, optim_args, clip_args=None): ...
300

301
# PyTorch optimizer wrappers
302
def Adam(optim_args, clip_args=None): ...
303
def SGD(optim_args, clip_args=None): ...
304
def RMSprop(optim_args, clip_args=None): ...
305
```
306

307
[Optimization](./optimization.md)
308

309
## Types
310

311
```python { .api }
312
from typing import Union, Optional, Dict, Any, Callable, Iterator, Sequence
313
from torch import Tensor
314
from pyro.distributions import Distribution, Transform, Constraint
315
from pyro.distributions.torch_distribution import TorchDistributionMixin
316
from pyro.params.param_store import ParamStoreDict
317
from pyro.poutine.runtime import InferDict
318
from pyro.poutine.plate_messenger import PlateMessenger
319

320
# Core types
321
ModelFunction = Callable[..., None]
322
GuideFunction = Callable[..., None] 
323
InitFunction = Callable[[str], Tensor]
324

325
# Distribution types
326
DistributionType = Union[Distribution, type]
327
ConstraintType = Union[Constraint, type]
328
TransformType = Union[Transform, type]
329

330
# Inference types  
331
OptimType = Union[torch.optim.Optimizer, pyro.optim.PyroOptim]
332
LossType = Union[pyro.infer.ELBO, Callable]
333

334
# Pyro-specific types
335
TorchDistributionMixin = pyro.distributions.torch_distribution.TorchDistributionMixin
336
ParamStoreDict = pyro.params.param_store.ParamStoreDict
337
InferDict = pyro.poutine.runtime.InferDict
338
PlateMessenger = pyro.poutine.plate_messenger.PlateMessenger
339
```