$ show pyro-ppl/pyro
Pyro is a deep probabilistic programming library built on PyTorch, providing a unified framework for Bayesian inference, variational inference, and Monte Carlo methods. Its architecture centers on composable effect handlers (poutine) that enable tracing, conditioning, replay, and reparameterization of probabilistic programs. The library includes a rich set of distributions, bijective transforms for normalizing flows, automatic guide construction, and a variety of inference algorithms (SVI, HMC, NUTS, SVGD, SMC). It also offers domain-specific contrib modules for Gaussian processes, time series forecasting, epidemiology, optimal experimental design, and multi-object tracking. While primarily a software library, Pyro has strong research roots, with many features inspired by recent papers and experimental modules that may evolve. The codebase is well-structured, extensively tested, and actively maintained, though some contrib modules carry instability warnings and a few internal patches to PyTorch may require attention.
The LLM identified 5 files worth highlighting for reuse:
Base class for Pyro's effect handler stack. Implements a composable context manager pattern where subclasses override `_pyro_sample`, `_pyro_param`, etc. to intercept and modify probabilistic program execution. This pattern is self-contained and can be extracted for other probabilistic programming or DSL projects.
Provides streaming statistics classes (CountMeanStats, CountMeanVarianceStats, StackStats) that support online computation and merging of summary statistics. The design is generic, depends only on PyTorch, and can be reused in any project requiring incremental or distributed statistics.
Implements Welford's online algorithm for covariance estimation, including a variant for arrowhead covariance matrices. The code is isolated, numerically stable, and directly reusable for adaptive MCMC or any online covariance computation.
Implements the Inverse Autoregressive Flow (IAF) transform, a key building block for normalizing flows. The module is self-contained, follows PyTorch's Transform interface, and demonstrates a reusable pattern for constructing autoregressive bijections.
Provides a custom backward pass for einsum operations used in sampling and marginalization. This module implements the adjoint method for einsum, enabling efficient gradient computation through discrete operations. It is isolated and can be adapted for other probabilistic or combinatorial optimization tasks.