Funsor-based Pyro

Effect handlers

class NamedMessenger(first_available_dim=None)[source]

Bases: pyro.poutine.reentrant_messenger.ReentrantMessenger

Base effect handler class for the :func:~`pyro.contrib.funsor.to_funsor` and :func:~`pyro.contrib.funsor.to_data` primitives. Any effect handlers that invoke these primitives internally or wrap code that does should inherit from NamedMessenger.

This design ensures that the global name-dim mapping is reset upon handler exit rather than potentially persisting until the entire program terminates.

class MarkovMessenger(history=1, keep=False)[source]

Bases: pyro.contrib.funsor.handlers.named_messenger.NamedMessenger

Handler for converting to/from funsors consistent with Pyro’s positional batch dimensions.

Parameters:
  • history (int) – The number of previous contexts visible from the current context. Defaults to 1. If zero, this is similar to pyro.plate.
  • keep (bool) – If true, frames are replayable. This is important when branching: if keep=True, neighboring branches at the same level can depend on each other; if keep=False, neighboring branches are independent (conditioned on their shared ancestors).
class GlobalNamedMessenger(first_available_dim=None)[source]

Bases: pyro.contrib.funsor.handlers.named_messenger.NamedMessenger

Base class for any new effect handlers that use the :func:~`pyro.contrib.funsor.to_funsor` and :func:~`pyro.contrib.funsor.to_data` primitives to allocate DimType.GLOBAL or DimType.VISIBLE dimensions.

Serves as a manual “scope” for dimensions that should not be recycled by :class:~`MarkovMessenger`: global dimensions will be considered active until the innermost GlobalNamedMessenger under which they were initially allocated exits.

to_funsor(x, output=None, dim_to_name=None, dim_type=<DimType.LOCAL: 0>)[source]
to_data(x, name_to_dim=None, dim_type=<DimType.LOCAL: 0>)[source]
class StackFrame(name_to_dim, dim_to_name, history=1, keep=False)[source]

Bases: object

Consistent bidirectional mapping between integer positional dimensions and names. Can be queried like a dictionary (value = frame[key], frame[key] = value).

class DimType[source]

Bases: enum.Enum

Enumerates the possible types of dimensions to allocate

LOCAL = 0
GLOBAL = 1
VISIBLE = 2
class DimRequest(value, dim_type)

Bases: tuple

dim_type

Alias for field number 1

value

Alias for field number 0

class DimStack[source]

Bases: object

Single piece of global state to keep track of the mapping between names and dimensions.

Replaces the plate DimAllocator, the enum EnumAllocator, the stack in MarkovMessenger, _param_dims and _value_dims in EnumMessenger, and dim_to_symbol in msg[‘infer’]

MAX_DIM = -25
DEFAULT_FIRST_DIM = -5
set_first_available_dim(dim)[source]
push_global(frame)[source]
pop_global()[source]
push_iter(frame)[source]
pop_iter()[source]
push_local(frame)[source]
pop_local()[source]
global_frame
local_frame
current_write_env
current_read_env

Collect all frames necessary to compute the full name <–> dim mapping and interpret Funsor inputs or batch shapes at any point in a computation.

allocate(key_to_value_request)[source]
names_from_batch_shape(batch_shape, dim_type=<DimType.LOCAL: 0>)[source]