Implementation:Pyro ppl Pyro TransformModule

Attribute	Value
Sources	pyro/distributions/torch_transform.py
Domains	Probabilistic Programming, Normalizing Flows, Learnable Transforms, Neural Networks
Last Updated	2026-02-09

Overview

Description

The TransformModule and ComposeTransformModule classes bridge PyTorch's distribution transforms system with its neural network module system. They enable transforms with learnable parameters (such as those used in normalizing flows) to participate in PyTorch's parameter management infrastructure.

TransformModule is a class that inherits from both torch.distributions.Transform and torch.nn.Module. By inheriting from nn.Module, any learnable parameters defined in subclasses are automatically registered and can be optimized by PyTorch optimizers, serialized/deserialized with state_dict(), and tracked by Pyro's parameter store when used inside PyroModule instances.

ComposeTransformModule similarly inherits from both torch.distributions.ComposeTransform and torch.nn.ModuleList. This allows a sequence of transform modules to be composed together while retaining proper parameter registration. Any nn.Module parts in the composition are automatically appended to the ModuleList, ensuring they appear in .parameters() and .named_modules() calls.

Both classes override __hash__ to use the default Python object identity hash (via super(torch.nn.Module, self).__hash__()), ensuring that transforms behave correctly when used as dictionary keys or in sets, which is important for caching mechanisms in PyTorch's transform system.

Usage

Any normalizing flow or learnable transform in Pyro should subclass TransformModule rather than the plain torch.distributions.Transform. When composing multiple such transforms, ComposeTransformModule should be used instead of torch.distributions.ComposeTransform to ensure all parameters are properly registered.

Code Reference

Source Location

Property	Value
File	pyro/distributions/torch_transform.py
Module	pyro.distributions.torch_transform
Repository	pyro-ppl/pyro

Signature

class TransformModule(torch.distributions.Transform, torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __hash__(self):
        return super(torch.nn.Module, self).__hash__()


class ComposeTransformModule(torch.distributions.ComposeTransform, torch.nn.ModuleList):
    def __init__(self, parts, cache_size=0):
        ...

    def __hash__(self):
        return super(torch.nn.Module, self).__hash__()

    def with_cache(self, cache_size=1):
        ...

Import

from pyro.distributions import TransformModule, ComposeTransformModule

# Or from the module directly:
from pyro.distributions.torch_transform import TransformModule, ComposeTransformModule

I/O Contract

TransformModule

Parameter	Type	Description
*args	varies	Passed to both torch.distributions.Transform and torch.nn.Module constructors via super().__init__()
**kwargs	varies	Passed to both parent constructors via super().__init__()

Inherited Interface	Source	Description
_call(x), _inverse(y)	Transform	Must be implemented by subclasses to define the forward and inverse mappings
log_abs_det_jacobian(x, y)	Transform	Must be implemented by subclasses for log probability computation
parameters(), named_parameters()	nn.Module	Enumerates all learnable parameters
state_dict(), load_state_dict()	nn.Module	Serialization of learnable parameters
to(device), cuda(), cpu()	nn.Module	Device transfer of parameters

ComposeTransformModule

Parameter	Type	Description
parts	list	List of transform instances. Any that are nn.Module instances are automatically registered in the ModuleList.
cache_size	int	Cache size for the composed transform. Default: 0.

Method	Return Type	Description
with_cache(cache_size=1)	ComposeTransformModule	Returns a new ComposeTransformModule with the specified cache size. Returns self if cache_size is unchanged.
__hash__()	int	Returns object identity hash (bypasses Module's hash override)

Usage Examples

Creating a Custom Learnable Transform

import torch
import torch.nn as nn
from pyro.distributions.torch_transform import TransformModule

class LearnableAffineTransform(TransformModule):
    """A simple learnable affine transform: y = scale * x + shift."""
    domain = torch.distributions.constraints.real
    codomain = torch.distributions.constraints.real
    bijective = True

    def __init__(self, dim):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(dim))
        self.shift = nn.Parameter(torch.zeros(dim))

    def _call(self, x):
        return self.scale * x + self.shift

    def _inverse(self, y):
        return (y - self.shift) / self.scale

    def log_abs_det_jacobian(self, x, y):
        return self.scale.abs().log().sum(-1)

# Parameters are automatically tracked
transform = LearnableAffineTransform(10)
print(list(transform.parameters()))  # [scale, shift]

Composing Multiple Learnable Transforms

import torch
from pyro.distributions import TransformModule, ComposeTransformModule, TransformedDistribution

# Assume we have custom TransformModule subclasses
# e.g., from pyro.distributions.transforms import Planar, Radial
from pyro.distributions.transforms import Planar

# Create a normalizing flow with multiple layers
transforms = [Planar(input_dim=5) for _ in range(3)]
flow = ComposeTransformModule(transforms)

# All parameters from all layers are accessible
total_params = sum(p.numel() for p in flow.parameters())
print(f"Total parameters: {total_params}")

# Use with a base distribution
base_dist = torch.distributions.Normal(torch.zeros(5), torch.ones(5))
flow_dist = TransformedDistribution(base_dist, [flow])

Using with PyroModule

import torch
import pyro
import pyro.distributions as dist
from pyro.nn import PyroModule

class FlowModel(PyroModule):
    def __init__(self):
        super().__init__()
        # TransformModule parameters are registered with Pyro's param store
        from pyro.distributions.transforms import Spline
        self.transform = Spline(input_dim=2, count_bins=8)

    def forward(self, data=None):
        base = dist.Normal(torch.zeros(2), torch.ones(2)).to_event(1)
        flow_dist = dist.TransformedDistribution(base, [self.transform])
        with pyro.plate("data", 100):
            return pyro.sample("x", flow_dist, obs=data)

Cache Control with ComposeTransformModule

from pyro.distributions.torch_transform import ComposeTransformModule
from pyro.distributions.transforms import Planar

# Build a composition of transforms
parts = [Planar(input_dim=3) for _ in range(5)]
composed = ComposeTransformModule(parts, cache_size=0)

# Enable caching for repeated forward/inverse calls
cached_composed = composed.with_cache(cache_size=1)
print(type(cached_composed))  # ComposeTransformModule

Related Pages

• TransformedDistribution -- Distribution formed by applying transforms to a base distribution
• Distributions_Init -- Central registry of all Pyro distributions including TransformModule
• Stable -- A distribution that benefits from reparameterization via transforms
• RelaxedStraightThrough -- Another distribution technique for differentiable discrete sampling