Implementation:Pyro ppl Pyro MarkovMessenger

Attribute	Value
File	pyro/poutine/markov_messenger.py
Module	pyro.poutine.markov_messenger
Lines	102
Parent Class	ReentrantMessenger (which extends Messenger)
Purpose	Declare Markov dependencies for memory-efficient sequential models
Architecture Role	Enables dimension recycling for enumeration of sequential (time-series) models
License	Apache-2.0 (Uber Technologies, Inc.)

Overview

MarkovMessenger declares Markov dependency structure in sequential models, serving as a statistical equivalent of a memory management arena. It tracks which sample sites are visible from the current context based on a configurable history window, enabling efficient dimension recycling during parallel enumeration.

The messenger supports three usage patterns:

• Context manager -- For wrapping blocks of code with Markov dependency.
• Iterator -- For iterating over Markov chains (e.g., time steps).
• Decorator -- For recursive functions.

Key parameters:

• history -- Number of previous contexts visible from the current one (default 1). If zero, behaves like pyro.plate.
• keep -- If True, frames are replayable so neighboring branches can depend on each other. If False, neighboring branches are independent.

The messenger annotates each sample site's infer dict with:

• _markov_scope -- A Counter of visible site names from the current and recent contexts.
• _markov_depth -- The nesting depth of Markov contexts.

Code Reference

class MarkovMessenger(ReentrantMessenger):
    def __init__(self, history=1, keep=False, dim=None, name=None):
        assert history >= 0
        self.history = history
        self.keep = keep
        self.dim = dim
        self.name = name
        self._iterable = None
        self._pos = -1
        self._stack = []
        super().__init__()

    def generator(self, iterable):
        self._iterable = iterable
        return self

    def __iter__(self):
        with ExitStack() as stack:
            for value in self._iterable:
                stack.enter_context(self)
                yield value

    def __enter__(self):
        self._pos += 1
        if len(self._stack) <= self._pos:
            self._stack.append(set())
        return super().__enter__()

    def __exit__(self, *args, **kwargs):
        if not self.keep:
            self._stack.pop()
        self._pos -= 1
        return super().__exit__(*args, **kwargs)

    def _pyro_sample(self, msg):
        if msg["done"] or type(msg["fn"]).__name__ == "_Subsample":
            return
        infer = msg["infer"]
        scope = infer.setdefault("_markov_scope", Counter())
        for pos in range(max(0, self._pos - self.history), self._pos + 1):
            scope.update(self._stack[pos])
        infer["_markov_depth"] = 1 + infer.get("_markov_depth", 0)
        self._stack[self._pos].add(msg["name"])

I/O Contract

Parameter	Type	Description
history	int (default 1)	Number of previous time steps visible from current context
keep	bool (default False)	Whether frames are kept for replayability across branches
dim	Optional[int]	Dimension for vectorized Markov (not yet implemented)
name	Optional[str]	Optional name for matching between model and guide (not yet implemented)

Message Effect	Description
msg["infer"]["_markov_scope"]	A Counter of site names visible from the current context (history window)
msg["infer"]["_markov_depth"]	Incremented nesting depth of Markov contexts

Usage Examples

Markov Chain as Iterator

def model(data):
    z = 0
    for i in pyro.markov(range(len(data))):
        z = pyro.sample(f"z_{i}", dist.Normal(z, 1))
        pyro.sample(f"x_{i}", dist.Normal(z, 0.5), obs=data[i])

Context Manager Usage

def model():
    with pyro.markov():
        z = pyro.sample("z", dist.Normal(0, 1))
    return z

Higher-Order Markov (history > 1)

# Second-order Markov chain: each step depends on 2 previous steps
for i in pyro.markov(range(T), history=2):
    z = pyro.sample(f"z_{i}", dist.Normal(0, 1))

Related Pages

• Pyro_ppl_Pyro_Messenger_Base -- Base handler protocol
• Pyro_ppl_Pyro_Poutine_Handlers -- The poutine.markov() factory function
• Pyro_ppl_Pyro_IndepMessenger -- Related: plate-based independence (Markov with history=0 is similar)
• Pyro_ppl_Pyro_Poutine_Runtime -- _EnumAllocator uses Markov scope for dimension recycling