Implementation:Pyro ppl Pyro MixtureOfDiagNormals

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

Overview

Description

MixtureOfDiagNormals implements a mixture of multivariate normal distributions where each component has a diagonal covariance matrix. The distribution is parameterized by K component means (a K x D matrix locs), K component scales (a K x D matrix coord_scale), and K mixture weight logits (a K-dimensional vector component_logits), where K is the number of mixture components and D is the dimensionality.

The key feature of this distribution is that it supports pathwise derivatives (reparameterized gradients) for samples, which is essential for efficient gradient-based variational inference. The implementation uses a custom torch.autograd.Function subclass (_MixDiagNormalSample) that implements the pathwise derivative estimator described in "Pathwise Derivatives for Multivariate Distributions" by Jankowiak & Karaletsos (arXiv:1806.01856).

The log_prob method uses a numerically stable computation via the log-sum-exp trick with a per-sample minimum subtraction for improved numerical precision. The backward pass of _MixDiagNormalSample computes analytic gradients for all three parameter groups (locs, coord_scale, and component_logits) using error functions and carefully constructed intermediate quantities.

Note: This distribution does not support dimension D = 1.

Usage

This distribution is used in variational inference scenarios where a mixture of Gaussians is needed as a flexible variational family or as a likelihood model. It is particularly suited for cases where reparameterized gradient estimates are required for stable training. The distribution supports both unbatched (K x D) and batched (... x B x K x D) parameter shapes.

Code Reference

Source Location

pyro/distributions/diag_normal_mixture.py

Signature

class MixtureOfDiagNormals(TorchDistribution):
    def __init__(self, locs, coord_scale, component_logits):
        ...

Import

from pyro.distributions import MixtureOfDiagNormals

I/O Contract

Inputs

Parameter	Type	Description
locs	torch.Tensor	K x D mean matrix (or ... x B x K x D for batched mode). Each row is the mean of one mixture component.
coord_scale	torch.Tensor	K x D scale matrix (or ... x B x K x D for batched mode). Each row contains the standard deviations for the diagonal covariance of one component. Must be positive.
component_logits	torch.Tensor	K-dimensional vector (or ... x B x K for batched mode) of unnormalized log mixture weights (softmax logits).

Outputs

Method	Return Type	Description
rsample(sample_shape)	torch.Tensor	Draws reparameterized samples with pathwise derivatives through the mixture. Returns tensors of shape sample_shape + batch_shape + (D,).
log_prob(value)	torch.Tensor	Computes the log probability of a value under the mixture, using numerically stable log-sum-exp.
expand(batch_shape)	MixtureOfDiagNormals	Returns a new distribution with expanded batch dimensions.

Usage Examples

import torch
from pyro.distributions import MixtureOfDiagNormals

K = 3  # number of mixture components
D = 5  # dimensionality

locs = torch.randn(K, D)
coord_scale = torch.ones(K, D)
component_logits = torch.zeros(K)

mix = MixtureOfDiagNormals(locs, coord_scale, component_logits)

# Draw reparameterized samples
samples = mix.rsample(sample_shape=(100,))
print("Sample shape:", samples.shape)  # (100, 5)

# Compute log probability
log_p = mix.log_prob(samples)
print("Log prob shape:", log_p.shape)  # (100,)

# Batched mode
batch_locs = torch.randn(4, K, D)
batch_scales = torch.ones(4, K, D)
batch_logits = torch.zeros(4, K)
batch_mix = MixtureOfDiagNormals(batch_locs, batch_scales, batch_logits)
batch_samples = batch_mix.rsample((10,))
print("Batched sample shape:", batch_samples.shape)  # (10, 4, 5)

Related Pages

• Pyro_ppl_Pyro_AVFMultivariateNormal -- Another distribution with custom pathwise gradient implementations
• Pyro_ppl_Pyro_Distribution_Base -- Base distribution class for all Pyro distributions
• Pyro_ppl_Pyro_Constraints -- Constraint definitions used for parameter validation