Implementation:Pyro ppl Pyro Infer Discrete

Metadata

Field	Value
Implementation ID	Pyro_ppl_Pyro_Infer_Discrete
Title	infer_discrete
Project	Pyro (pyro-ppl/pyro)
File	pyro/infer/discrete.py, Lines 181-231
Implements	Pyro_ppl_Pyro_Discrete_Posterior_Decoding
Repository	https://github.com/pyro-ppl/pyro

Summary

infer_discrete is a Pyro poutine (effect handler) that samples discrete sites marked with site["infer"]["enumerate"] = "parallel" from the posterior, conditioned on observations. It supports both posterior sampling (temperature=1, forward-filter backward-sample) and MAP decoding (temperature=0, Viterbi-like).

Signature

def infer_discrete(
    fn=None,
    first_available_dim=None,
    temperature=1,
    *,
    strict_enumeration_warning=True
)

Import

from pyro.infer import infer_discrete

Parameters

Parameter	Type	Default	Description
fn	callable or None	None	A stochastic function (model) containing Pyro primitive calls with discrete enumerated sites. If None, returns a decorator.
first_available_dim	int	None	The first tensor dimension (counting from the right, as a negative integer) available for parallel enumeration. This dimension and all dimensions to the left may be used internally by Pyro for enumeration.
temperature	int	1	Either 1 (sample from posterior via forward-filter backward-sample) or 0 (MAP/Viterbi decoding via max-product).
strict_enumeration_warning	bool	True	Whether to warn if no enumerated sample sites are found. Keyword-only argument.

Returns

Type	Description
callable	A functools.partial wrapping _sample_posterior. When called with the model's arguments, it returns the model's return value with discrete sites replaced by posterior samples or MAP values.

Usage Patterns

As a Decorator

@infer_discrete(first_available_dim=-1, temperature=0)
@config_enumerate
def viterbi_decoder(data, hidden_dim=10):
    transition = 0.3 / hidden_dim + 0.7 * torch.eye(hidden_dim)
    means = torch.arange(float(hidden_dim))
    states = [0]
    for t in pyro.markov(range(len(data))):
        states.append(
            pyro.sample("states_{}".format(t),
                        dist.Categorical(transition[states[-1]]))
        )
        pyro.sample("obs_{}".format(t),
                     dist.Normal(means[states[-1]], 1.0),
                     obs=data[t])
    return states  # returns maximum likelihood states

map_states = viterbi_decoder(data)

As a Function

model_fn = config_enumerate(model)
decoder = infer_discrete(model_fn, first_available_dim=-1, temperature=1)
posterior_sample = decoder(data)

After SVI Training (with guide replay)

# After training with SVI + TraceEnum_ELBO
guide_trace = poutine.trace(trained_guide).get_trace(data)
replayed_model = poutine.replay(model, trace=guide_trace)
decoder = infer_discrete(replayed_model, first_available_dim=-2, temperature=1)
decoded = decoder(data)

Internal Mechanism

The function (lines 181-231) supports both decorator and function call patterns:

def infer_discrete(fn=None, first_available_dim=None, temperature=1, *,
                   strict_enumeration_warning=True):
    assert first_available_dim < 0, first_available_dim
    if fn is None:  # support use as a decorator
        return functools.partial(
            infer_discrete,
            first_available_dim=first_available_dim,
            temperature=temperature,
        )
    return functools.partial(
        _sample_posterior,
        fn,
        first_available_dim,
        temperature,
        strict_enumeration_warning,
    )

The core computation is performed by _sample_posterior (lines 41-55) and _sample_posterior_from_trace (lines 58-178):

1. Create enumerated trace: Runs the model under EnumMessenger(first_available_dim) to create a trace where discrete sites are expanded over all possible values.
2. Collect log-probability terms: Organizes masked log-probabilities into cost_terms and enum_terms by their plate ordinal.
3. Run forward-backward algorithm: Calls contract_tensor_tree with the appropriate ring:
   • SampleRing (temperature=1): sum-product with stochastic backward sampling
   • MapRing (temperature=0): max-product with argmax backward tracing
4. Construct collapsed trace: Gathers the decoded discrete values and adjusts cond_indep_stack for each site.
5. Replay model: Runs the model conditioned on the collapsed trace via SamplePosteriorMessenger.

Ring Selection

Temperature	Ring	Forward Operation	Backward Operation
1	SampleRing	Log-sum-exp (marginals)	Categorical sampling from posterior
0	MapRing	Log-max (Viterbi)	Argmax traceback

Warning

The log_prob values in the inferred model's trace are not meaningful and may change in future releases. They should not be used for downstream computation.

Complete Example

import torch
import pyro
import pyro.distributions as dist
from pyro.infer import config_enumerate, infer_discrete

# 3-state HMM
hidden_dim = 3
transition_probs = torch.tensor([
    [0.7, 0.2, 0.1],
    [0.1, 0.8, 0.1],
    [0.2, 0.1, 0.7],
])
emission_locs = torch.tensor([0.0, 5.0, 10.0])

@infer_discrete(first_available_dim=-1, temperature=0)
@config_enumerate
def viterbi_model(data):
    states = [0]
    for t in pyro.markov(range(len(data))):
        state = pyro.sample(
            "state_{}".format(t),
            dist.Categorical(transition_probs[states[-1]])
        )
        states.append(state)
        pyro.sample(
            "obs_{}".format(t),
            dist.Normal(emission_locs[state], 1.0),
            obs=data[t]
        )
    return states

# Generate synthetic data and decode
data = torch.tensor([0.1, 0.3, 5.2, 4.8, 9.5, 10.1, 5.0])
map_states = viterbi_model(data)
print("MAP states:", [s.item() for s in map_states[1:]])
# Expected: states near [0, 0, 1, 1, 2, 2, 1]

Related Pages

Implements Principle

• Principle:Pyro_ppl_Pyro_Discrete_Posterior_Decoding

Related Implementations

• Pyro_ppl_Pyro_Config_Enumerate -- Required to annotate discrete sites for enumeration before calling infer_discrete.
• Pyro_ppl_Pyro_Contract_Tensor_Tree -- The core tensor contraction algorithm used during the forward pass of decoding.
• Pyro_ppl_Pyro_Pyro_Markov -- Declares Markov structure for memory-efficient enumeration in sequential models.