Sampling Observers

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

init() -> PyTree

Initialize the memory for the observer. Defaults to None.

thrml.StateObserver(thrml.AbstractObserver)

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

__init__(blocks_to_sample: list[thrml.Block])

Initialize self. See help(type(self)) for accurate signature.

thrml.MomentAccumulatorObserver(thrml.AbstractObserver)

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.

__init__(moment_spec: Sequence[Sequence[Sequence[thrml.AbstractNode]]], f_transform: typing.Callable = <function _f_identity>)

Create a MomentAccumulatorObserver.

Arguments:

• moment_spec: A 3 depth sequence. The first is a sequence over different moment types. A given moment type should have the same number of nodes in each moment. Then for each moment type, there is a sequence over moments. Each given moment is defined by a certain set of nodes.

  For example, to get the first and second moments on a simple o-o graph, it would be

  [ [(node1,), (node2,)], [(node1, node2)] ] - f_transform: A function that takes in (state, blocks) and returns something with the same structure as state. This is used to apply functions to the samples before moments are computed. i.e this function defines a transformation of the state variable \(y=f(x)\), such that the accumulated moments are of the form \(\langle f(x_1) f(x_2) \rangle\).