Implementation:Pyro ppl Pyro ZukoToPyro

Property	Value
Module	pyro.contrib.zuko
Source	pyro/contrib/zuko.py
Lines	82
Classes	ZukoToPyro
Parent Class	pyro.distributions.TorchDistribution
Dependencies	torch, pyro

Overview

ZukoToPyro wraps a Zuko distribution (or any torch.distributions.Distribution) as a Pyro-compatible distribution, enabling seamless integration of Zuko-based normalizing flows within Pyro's inference pipelines.

Zuko is a separate library for normalizing flows that provides high-quality flow implementations. The wrapper enables:

• Using Zuko flows as guides in Pyro SVI.
• Using Zuko flows as priors in Pyro models.
• Efficient sampling + scoring via Zuko's rsample_and_log_prob method when available.

A key optimization is that when the wrapped distribution has an rsample_and_log_prob method (as Zuko flows do), the log probability is cached during sampling and reused during scoring, avoiding redundant computation.

Code Reference

Class: ZukoToPyro

Constructor:

• dist (torch.distributions.Distribution): The distribution to wrap.

Properties:

• has_rsample: Delegates to the wrapped distribution.
• event_shape: Delegates to the wrapped distribution.
• batch_shape: Delegates to the wrapped distribution.

Methods:

• __call__(shape=()): Draws a sample. If the distribution has rsample_and_log_prob, uses it and caches the log probability keyed by the sample tensor identity. Otherwise falls back to rsample or sample.

• log_prob(x): Returns the log probability. First checks the cache (populated by __call__); if the sample is not found, computes directly via the wrapped distribution.

• expand(*args, **kwargs): Returns a new ZukoToPyro wrapping the expanded distribution.

I/O Contract

Method	Input	Output
__init__	dist: torch.distributions.Distribution	ZukoToPyro instance
__call__(shape)	shape: torch.Size	Tensor (sample)
log_prob(x)	x: Tensor	Tensor (log probability)
expand(*args)	Same as Distribution.expand	ZukoToPyro

Usage Examples

import torch
import pyro
import pyro.distributions as dist
from pyro.contrib.zuko import ZukoToPyro

# Wrap a Zuko flow as a Pyro distribution
# (requires zuko to be installed: pip install zuko)
import zuko

flow = zuko.flows.MAF(features=5)

# flow() creates a torch.distributions.Distribution
zuko_dist = flow()

# Wrap as Pyro distribution
pyro_dist = ZukoToPyro(zuko_dist)

# Use in Pyro model
x = pyro_dist((2, 3))  # sample with shape (2, 3, 5)
log_p = pyro_dist.log_prob(x)  # uses cached value if available

# Use as a flow guide in SVI
def guide(data):
    flow = zuko.flows.MAF(features=data.shape[-1])
    pyro.sample("z", ZukoToPyro(flow()))

# Use with pyro.plate
with pyro.plate("data", 42):
    z = pyro.sample("z", ZukoToPyro(flow()))

Related Pages

• Pyro_ppl_Pyro_MiniPyro -- Simplified Pyro for pedagogical reference
• Pyro_ppl_Pyro_AutoRegressiveNN -- MADE networks used in Pyro's own autoregressive flows