Funsor-based Pyro

plate(*args, **kwargs)

Effect handlers

enum(fn=None, *args, **kwargs)

Convenient wrapper of EnumMessenger

This version of EnumMessenger uses to_data to allocate a fresh enumeration dim for each discrete sample site.

markov(fn=None, *args, **kwargs)

Convenient wrapper of MarkovMessenger

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).

named(fn=None, *args, **kwargs)

Convenient wrapper of NamedMessenger

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.

plate(fn=None, *args, **kwargs)

Convenient wrapper of PlateMessenger

Combines new IndepMessenger implementation with existing BroadcastMessenger. Should eventually be a drop-in replacement for pyro.plate.

replay(fn=None, *args, **kwargs)

Convenient wrapper of ReplayMessenger

This version of ReplayMessenger is almost identical to the original version, except that it calls to_data on the replayed funsor values. This may result in different unpacked shapes, but should produce correct allocations.

trace(fn=None, *args, **kwargs)

Convenient wrapper of TraceMessenger

Setting pack_online=True packs online instead of after the fact, converting all distributions and values to Funsors as soon as they are available.

Setting pack_online=False computes information necessary to do packing after execution. Each sample site is annotated with a “dim_to_name” dictionary, which can be passed directly to funsor.to_funsor.

class NamedMessenger(first_available_dim=None)

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)

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)

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>)

to_data(x, name_to_dim=None, dim_type=<DimType.LOCAL: 0>)

class StackFrame(name_to_dim, dim_to_name, history=1, keep=False)

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

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

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)

push_global(frame)

pop_global()

push_iter(frame)

pop_iter()

push_local(frame)

pop_local()

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)

names_from_batch_shape(batch_shape, dim_type=<DimType.LOCAL: 0>)