Implementation:Pyro ppl Pyro ConditionalDistribution

Knowledge Sources	Pyro
Domains	Probability_Distributions
Last Updated	2026-02-09 09:00 GMT

Overview

Description

The conditional module provides a family of abstract base classes and concrete utilities for defining conditional distributions and conditional transforms in Pyro. A conditional distribution is a distribution whose parameters depend on an external context tensor. The module is the foundation for building conditional normalizing flows and other context-dependent probabilistic models.

The module defines the following key classes:

• ConditionalDistribution -- An abstract base class (ABC) with a single abstract method condition(context) that returns a torch.distributions.Distribution.
• ConditionalTransform -- An abstract base class with a condition(context) method that returns a torch.distributions.Transform.
• ConditionalTransformModule -- Extends both ConditionalTransform and torch.nn.Module so that conditional transforms with learnable parameters (such as normalizing flows) can participate in PyTorch's parameter management.
• ConditionalComposeTransformModule -- A conditional analogue of ComposeTransformModule that composes a sequence of conditional transforms. It extends both ConditionalTransformModule and torch.nn.ModuleList.
• ConditionalTransformedDistribution -- Combines a base distribution (optionally conditional) with a list of conditional transforms. When condition(context) is called, it conditions both the base distribution and all transforms, returning a standard TransformedDistribution.
• ConstantConditionalDistribution and ConstantConditionalTransform -- Wrapper classes that adapt non-conditional distributions and transforms to the conditional interface by ignoring the context.

Usage

This module is used to build context-dependent distribution families, particularly conditional normalizing flows. Users subclass ConditionalTransformModule to implement learnable conditional transforms, compose them via ConditionalComposeTransformModule, and wrap everything in a ConditionalTransformedDistribution. At inference time, calling .condition(context) yields a standard PyTorch distribution that can be used for sampling and log-probability evaluation.

Code Reference

Source Location

pyro/distributions/conditional.py

Signature

class ConditionalDistribution(ABC):
    def condition(self, context): ...

class ConditionalTransform(ABC):
    def condition(self, context): ...

class ConditionalTransformModule(ConditionalTransform, torch.nn.Module):
    ...

class ConditionalComposeTransformModule(ConditionalTransformModule, torch.nn.ModuleList):
    def __init__(self, transforms, cache_size=0): ...
    def condition(self, context): ...

class ConditionalTransformedDistribution(ConditionalDistribution):
    def __init__(self, base_dist, transforms): ...
    def condition(self, context): ...

Import

from pyro.distributions.conditional import (
    ConditionalDistribution,
    ConditionalTransform,
    ConditionalTransformModule,
    ConditionalComposeTransformModule,
    ConditionalTransformedDistribution,
)

I/O Contract

Inputs

Class	Parameter	Type	Description
ConditionalDistribution	context	torch.Tensor	Context tensor passed to the condition method to produce a concrete distribution.
ConditionalComposeTransformModule	transforms	list	A list of ConditionalTransform or torch.distributions.Transform objects to compose.
ConditionalComposeTransformModule	cache_size	int	Cache size for the composed transform (must be 0 or 1).
ConditionalTransformedDistribution	base_dist	Distribution or ConditionalDistribution	The base distribution, optionally conditional.
ConditionalTransformedDistribution	transforms	list	A list of conditional or unconditional transforms to apply.

Outputs

Method	Return Type	Description
ConditionalDistribution.condition(context)	torch.distributions.Distribution	A concrete distribution conditioned on the given context.
ConditionalTransform.condition(context)	torch.distributions.Transform	A concrete transform conditioned on the given context.
ConditionalComposeTransformModule.condition(context)	ComposeTransformModule	A composed transform with each component conditioned on the context.
ConditionalTransformedDistribution.condition(context)	TransformedDistribution	A TransformedDistribution with base distribution and transforms conditioned on the context.

Usage Examples

import torch
import pyro.distributions as dist
from pyro.distributions.conditional import (
    ConditionalComposeTransformModule,
    ConditionalTransformedDistribution,
)

# Define a conditional flow stack
class ConditionalFlowStack(ConditionalComposeTransformModule):
    def __init__(self, input_dim, context_dim, hidden_dims, num_flows):
        super().__init__([
            dist.transforms.conditional_planar(input_dim, context_dim, hidden_dims)
            for _ in range(num_flows)
        ], cache_size=1)

# Build a conditional transformed distribution
cond_dist = ConditionalTransformedDistribution(
    dist.Normal(torch.zeros(3), torch.ones(3)).to_event(1),
    [ConditionalFlowStack(3, 2, [8, 8], num_flows=4).inv]
)

# Condition on context and compute negative log-likelihood
context = torch.rand(10, 2)
data = torch.rand(10, 3)
nll = -cond_dist.condition(context).log_prob(data)

Related Pages

• Pyro_ppl_Pyro_Distribution_Base -- Base distribution class for all Pyro distributions
• Pyro_ppl_Pyro_FoldedDistribution -- A transformed distribution implementation
• Pyro_ppl_Pyro_Constraints -- Constraint definitions used by distributions