Principle:Pyro ppl Pyro Normalizing Flow Integration

Knowledge Sources	Normalizing Flows for Probabilistic Modeling and Inference Variational Inference with Normalizing Flows Zuko: Normalizing Flows in PyTorch
Domains	Normalizing Flows, Density Estimation, Variational Inference
Last Updated	2026-02-09 09:00 GMT

Overview

Normalizing flow integration provides adapters that convert normalizing flow distributions from external libraries into Pyro-compatible distribution objects, enabling their use as flexible variational approximations within Pyro's inference framework.

Description

Normalizing flows are a class of generative models that construct complex distributions by applying a sequence of invertible, differentiable transformations to a simple base distribution (typically a standard Gaussian). The key property is that the density of the transformed distribution can be computed exactly using the change-of-variables formula.

While Pyro has its own normalizing flow implementations, the broader ecosystem includes specialized libraries (such as Zuko) that offer additional flow architectures, optimized implementations, or research-oriented features. The normalizing flow integration layer provides adapters that bridge these external libraries with Pyro's distribution system.

The adapter pattern involves:

1. Wrapping the external flow's distribution object so it implements Pyro's TorchDistribution interface.
2. Delegating sample(), log_prob(), and rsample() calls to the underlying flow.
3. Translating shape conventions: different libraries may have different conventions for batch vs. event dimensions.
4. Ensuring compatibility with Pyro's effect handlers, plates, and inference algorithms.

This design follows the adapter pattern from software engineering: it allows Pyro to leverage innovations in the normalizing flow community without coupling to any specific external implementation.

Usage

Use normalizing flow integration when:

• You want to use a flow architecture from an external library (e.g., Zuko) as a Pyro guide.
• The external library offers optimized or specialized flows not available in Pyro natively.
• Building flexible variational approximations that go beyond mean-field or low-rank Gaussian guides.
• Performing density estimation with state-of-the-art flow architectures within a Pyro model.

Theoretical Basis

Change of variables formula:

# Given: base distribution z ~ p_0(z) (e.g., standard Normal)
# Transform: x = f(z) where f is invertible and differentiable
# Inverse: z = f^{-1}(x)

# Density of x:
# p(x) = p_0(f^{-1}(x)) * |det(df^{-1}/dx)|
#       = p_0(f^{-1}(x)) * |det(df/dz)|^{-1}

# Log-density:
# log p(x) = log p_0(f^{-1}(x)) - log |det(df/dz)|_{z=f^{-1}(x)}

Composition of flows:

# Chain K transformations: x = f_K(f_{K-1}(...f_1(z)...))
# Log-density:
# log p(x) = log p_0(z) - sum_{k=1}^{K} log |det(df_k/dz_k)|

# where z_0 = z, z_k = f_k(z_{k-1}), z_K = x

# Each f_k is designed so that det(df_k/dz_k) is cheap to compute:
# - Triangular Jacobian: O(d) determinant
# - Low-rank perturbation: O(d) determinant via matrix determinant lemma

Adapter pattern for integration:

# External library provides:
class ExternalFlow:
    def sample(n_samples) -> Tensor
    def log_prob(x) -> Tensor

# Adapter wraps it for Pyro:
class FlowToPyro(TorchDistribution):
    def __init__(self, external_flow):
        self.flow = external_flow
        # Set batch_shape, event_shape from flow

    def rsample(self, sample_shape):
        return self.flow.sample(sample_shape)

    def log_prob(self, value):
        return self.flow.log_prob(value)

    @property
    def has_rsample(self):
        return True  # flows are reparameterizable by construction

# Usage in Pyro:
# guide = FlowToPyro(zuko.flows.MAF(features=d, transforms=5))
# svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

Related Pages

• Implementation:Pyro_ppl_Pyro_ZukoToPyro