Implementation:Pyro ppl Pyro Poutine Condition

Metadata

Field	Value
Page Type	Implementation (API Doc)
Knowledge Sources	Repo (Pyro)
Domains	Bayesian_Inference, Probabilistic_Programming
Last Updated	2026-02-09 12:00 GMT

Overview

Concrete effect handler for programmatically conditioning a Pyro model on observed data by intercepting sample sites and replacing their values with provided observations.

Description

ConditionMessenger is Pyro's effect handler (messenger) that implements programmatic observation conditioning. It is exposed to users through the convenience function poutine.condition(fn, data), which wraps a model callable so that specified sample sites are treated as observed.

When the conditioned model is executed:

1. Each pyro.sample statement generates a message in Pyro's effect handling system.
2. The ConditionMessenger intercepts each message during the _pyro_sample method.
3. If the sample site name matches a key in the data dictionary, the messenger:
   • Sets the message's value to the corresponding tensor from data.
   • Sets is_observed = True on the message.
4. The downstream inference machinery then scores the observed value under the distribution (computes the log probability) rather than drawing a random sample.

The data argument can be either:

• A dictionary (Dict[str, torch.Tensor]) mapping sample site names to observed tensor values.
• A Trace object from a previous model execution, in which case the values are extracted from the trace's sample sites.

ConditionMessenger inherits from Messenger, Pyro's base class for effect handlers. It can be composed with other messengers (e.g., trace, replay, block) using Pyro's handler stacking mechanism.

Code Reference

Source Location

Pyro repo, file: pyro/poutine/condition_messenger.py, lines L15-71.

Factory Function Location

poutine.condition is defined in pyro/poutine/handlers.py.

Signature

# Factory function (user-facing API)
def condition(fn, data):
    """
    :param fn: a stochastic function (callable containing Pyro primitive calls)
    :param data: a dict or Trace mapping sample site names to observed values
    :returns: stochastic function wrapped in a ConditionMessenger
    """
    ...

# Underlying messenger class
class ConditionMessenger(Messenger):
    def __init__(self, data):
        super().__init__()
        self.data = data

    def _pyro_sample(self, msg):
        ...

Import

import pyro.poutine as poutine

# Use via the factory function
conditioned_model = poutine.condition(model, data={"obs": observed_data})

I/O Contract

Inputs

Parameter	Type	Required	Description
fn	callable	Yes	A Pyro model (stochastic function containing pyro.sample calls). This is the model to be conditioned.
data	Union[Dict[str, torch.Tensor], Trace]	Yes	A dictionary mapping sample site names to observed tensor values, or a Trace object from a previous model execution. Sites named in this dict will be treated as observed.

Outputs

Output	Type	Description
Conditioned model	callable	A new callable that behaves like the original model but with the specified sample sites treated as observed. When called, the conditioned model scores the provided observed values under their respective distributions instead of sampling.

Usage Examples

Basic Conditioning with a Dictionary

import torch
import pyro
import pyro.distributions as dist
import pyro.poutine as poutine

def model():
    mu = pyro.sample("mu", dist.Normal(0, 10))
    sigma = pyro.sample("sigma", dist.HalfNormal(10))
    obs = pyro.sample("obs", dist.Normal(mu, sigma))
    return obs

# Condition the model on observed data
data = torch.tensor(5.0)
conditioned_model = poutine.condition(model, data={"obs": data})

# Now "obs" is treated as observed in any inference algorithm

Using poutine.condition with MCMC

import torch
import pyro
import pyro.distributions as dist
import pyro.poutine as poutine
from pyro.infer.mcmc import NUTS, MCMC

def model(N):
    mu = pyro.sample("mu", dist.Normal(0, 10))
    sigma = pyro.sample("sigma", dist.HalfNormal(10))
    with pyro.plate("data", N):
        z = pyro.sample("z", dist.Normal(mu, sigma))
    return z

# Observed data
observed_z = torch.randn(100) * 2 + 5

# Condition the model programmatically
conditioned = poutine.condition(model, data={"z": observed_z})

# Run MCMC on the conditioned model
nuts_kernel = NUTS(conditioned)
mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=500)
mcmc.run(100)  # pass N=100

samples = mcmc.get_samples()
print("Posterior mu:", samples["mu"].mean().item())
print("Posterior sigma:", samples["sigma"].mean().item())

Conditioning on Multiple Sites

import torch
import pyro
import pyro.distributions as dist
import pyro.poutine as poutine

def model():
    z1 = pyro.sample("z1", dist.Normal(0, 1))
    z2 = pyro.sample("z2", dist.Normal(z1, 1))
    x = pyro.sample("x", dist.Normal(z2, 0.1))
    return x

# Condition on both z2 and x
conditioned = poutine.condition(model, data={
    "z2": torch.tensor(3.0),
    "x": torch.tensor(2.9),
})

Composing with Other Handlers

import pyro.poutine as poutine

# Condition and trace simultaneously
conditioned = poutine.condition(model, data={"obs": data})
traced = poutine.trace(conditioned)
trace = traced.get_trace()

# Inspect the trace to verify conditioning
assert trace.nodes["obs"]["is_observed"] == True
print("Log prob of obs:", trace.nodes["obs"]["log_prob"].item())

Related Pages

• Principle:Pyro_ppl_Pyro_Observation_Conditioning