THRML

Sampling observers

Observers accumulate statistics over a chain as it runs, so you read off moments or stored states without materializing every sample.

AbstractObserverclass
AbstractObserver()

Interface for objects that inspect the sampling program while it is running.

A concrete Observer is called once per block-sampling iteration and can maintain an arbitrary "carry" state across calls (e.g. running averages, histogram buffers, log-probs, etc.).

__call__method
__call__(program: 'BlockSamplingProgram', state_free: list[PyTree[Array]], state_clamped: list[PyTree[Array]], carry: ~ObserveCarry, iteration: Int[Array, '']) -> tuple[~ObserveCarry, PyTree]

Make an observation.

This function is called at the end of a block-sampling iteration and can record information about the current state of the sampling program that might be useful for something later.

Arguments:

  • program: The sampling program that is running when this function is called.
  • state_free: The current state of the free nodes involved in the sampling program.
  • state_clamped: The state of the clamped nodes involved in the sampling program.
  • carry: The "memory" available to this observer. This function should modify this PyTree to record information about the sampling program.
  • iteration: How many iterations of block sampling have happened before this function was called.

Returns:

A tuple, where the first element is the updated carry, and the second is a PyTree that will be recorded by the sampler.

initmethod
init() -> PyTree

Initialize the memory for the observer. Defaults to None.

StateObserverclass
StateObserver(blocks_to_sample: list[Block])

Observer which logs the raw state of some set of nodes.

Attributes:

  • blocks_to_sample: the list of Blocks which the states are logged for
blocks_to_sampleattribute
blocks_to_sample: list[Block]
initmethod
init() -> PyTree

Initialize the memory for the observer. Defaults to None.

__call__method
__call__(program: 'BlockSamplingProgram', state_free: list['_State'], state_clamped: list['_State'], carry: None, iteration: Int[Array, '']) -> tuple[None, PyTree]

Simply returns the state of the blocks that are being logged to be recorded by the sampler.

MomentAccumulatorObserverclass
MomentAccumulatorObserver(moment_spec: Sequence[Sequence[Sequence[AbstractNode]]], f_transform: Callable = _f_identity)

Observer that accumulates and updates the provided moments.

It doesn't log any samples, and will only accumulate moments. Note that this observer does not scale the accumulated values by the number of times it was called. It simply records a running sum of a product of some state variables,

$$\sum_i f(x_1^i) f(x_2^i) \dots f(x_N^i)$$

Attributes:

  • blocks_to_sample: the blocks to accumulate the moments over. These are for constructing the final state, and aren't truly "blocks" in the algorithmic sense (they can be connected to each other). There is one block per node type.
  • flat_nodes_list: a list of all of the nodes in the moments (each occurring only once, so len(set(x)) = len(x)).
  • flat_to_type_slices_list: a list over node types in which each element is an array of indices of the flat_node_list which that type corresponds to
  • flat_to_full_moment_slices: a list over moment types in which each element is a 2D array, which matches the shape of the moment_spec[i] and of which each element is the index in the flat_node_list.
  • f_transform: the element-wise transformation $f$ to apply to sample values before accumulation.
blocks_to_sampleattribute
blocks_to_sample: list[Block]
flat_nodes_listattribute
flat_nodes_list: list[AbstractNode]
flat_to_type_slices_listattribute
flat_to_type_slices_list: list[Int[Array, 'nodes_in_slice']]
flat_to_full_moment_slicesattribute
flat_to_full_moment_slices: list[Int[Array, 'num_groups nodes_in_moment']]
f_transformattribute
f_transform: Callable
__call__method
__call__(program: 'BlockSamplingProgram', state_free: list[PyTree[Array]], state_clamped: list[PyTree[Array]], carry: list[Array], iteration: Int[Array, '']) -> tuple[list[Array], PyTree]

Accumulate the moments via carry. Does not return anything for the sampler to write down.

initmethod
init() -> list[Array]

Initialize the memory that will store the accumulated values.