tfp.experimental.mcmc.ExpectationsReducer

Reducer that computes a running expectation.

Inherits From: Reducer

tfp.experimental.mcmc.ExpectationsReducer(
    transform_fn=_get_sample, name=None
)

ExpectationsReducer calculates expectations over some arbitrary structure of transform_fns. A transform_fn is a function that accepts a Markov chain sample and kernel results, and outputs the relevant value for expectation calculation. In other words, if we denote a transform_fn as f(x,y), ExpectationsReducer approximates E[f(x,y)] for all provided functions. The finalized expectation will also identically mimic the structure of transform_fn.

As with all reducers, ExpectationsReducer does not hold state information; rather, it stores supplied metadata. Intermediate calculations are held in a state object, which is returned via initialize and one_step method calls.

transform_fn A (possibly nested) structure of functions that accept a chain state and kernel results. Defaults to simply returning the incoming chain state.
name Python str name prefixed to Ops created by this function. Default value: None (i.e., 'expectations_reducer').

name

parameters

transform_fn

Methods

finalize

View source

finalize(
    final_reducer_state
)

Finalizes expectation calculation from the final_reducer_state.

If the finalized method is invoked on a stream of no inputs, a corresponding structure of tf.ones will be returned.

Args
final_reducer_state ExpectationsReducerState that represents the final reducer state.

Returns
expectation an estimate of the expectation with identical structure to self.transform_fn.

initialize

View source

initialize(
    initial_chain_state, initial_kernel_results=None
)

Initializes an empty ExpectationsReducerState.

Args
initial_chain_state A (possibly nested) structure of Tensors or Python lists of Tensors representing the current state(s) of the Markov chain(s).
initial_kernel_results A (possibly nested) structure of Tensors representing internal calculations made in a related TransitionKernel.

Returns
state ExpectationsReducerState representing a stream of no inputs.

one_step

View source

one_step(
    new_chain_state,
    current_reducer_state,
    previous_kernel_results=None,
    axis=None
)

Update the current_reducer_state with a new chain state.

Chunking semantics are specified by the axis parameter. If chunking is enabled (axis is not None), all elements along the specified axis will be treated as separate samples. If a single scalar value is provided for a non-scalar sample structure, that value will be used for all elements in the structure. If not, an identical structure must be provided.

Args
new_chain_state A (possibly nested) structure of incoming chain state(s) with shape and dtype compatible with those used to initialize the current_reducer_state.
current_reducer_state ExpectationsReducerState representing the current reducer state.
previous_kernel_results A (possibly nested) structure of Tensors representing internal calculations made in a related TransitionKernel.
axis If chunking is desired, this is a (possibly nested) structure of integers that specifies the axis with chunked samples. For individual samples, set this to None. By default, samples are not chunked (axis is None).

Returns
new_reducer_state ExpectationsReducerState with updated running statistics. It tracks a running total and the number of processed samples.