tfp.experimental.mcmc.CovarianceReducer

Reducer that computes a running covariance.

Inherits From: Reducer

tfp.experimental.mcmc.CovarianceReducer(
    event_ndims=None, transform_fn=_get_sample, ddof=0, name=None
)

By default, CovarianceReducer computes covariance directly over the latent state, ignoring kernel results; however, it can also compute a running covariance on samples evaluated on some arbitrary structure of transform_fns. A trasform_fn is defined as any function that accepts the chain state and kernel results, and ouptuts a Tensor or nested structure of Tensors to compute the covariance over. To be explicit, if we denote a transform_fn as f(x, y), then CovarianceReducer will compute Cov(f(x, y)). If some structure of transform_fns are given, the final statistic (as returned by the finalize method) will mimic the results of that structure.

As with all reducers, CovarianceReducer 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.

CovarianceReducer is meant to fit into the larger Streaming MCMC framework. RunningCovariance in tfp.experimental.stats is better suited for more generic streaming covariance needs.

There are two ways to stream over MCMC samples. The first is to manually wrap CovarianceReducer in a tfp.experimental.mcmc.WithReductions TransitionKernel. Doing so enables a wide variety of possible compositions. For example, to perform covariance calculation on samples that survive thinning and burn-in, WithReductions should wrap around an appropriate SampleDiscardingKernel.

If the sole objective is to estimate covariance, it may be more convenient to use a pre-defined driver like tfp.experimental.mcmc.sample_fold. sample_fold will automatically create the Transition Kernel onion, apply kernel samples to update reducer states, and finalize calculations.

python

kernel = ...
reducer = tfp.experimental.mcmc.CovarianceReducer()
covariance_estimate, _, _ = tfp.experimental.mcmc.sample_fold(
    num_steps=...,
    current_state=...,
    previous_kernel_results=...,
    kernel=kernel,
    reducers=reducer,
)
covariance_estimate  # estimate of covariance (not `CovarianceReducerState`)

event_ndims A (possibly nested) structure of integers, or None. Defines the number of inner-most dimensions that represent the event shape. Specifying None returns all cross product terms (no batching) and is the default.
transform_fn A (possibly nested) structure of functions to evaluate the incoming chain states on before covariance calculation. The default is a single function that returns the chain state.
ddof A (possibly nested) structure of integers that represent the requested dynamic degrees of freedom. For example, use ddof=0 for population covariance and ddof=1 for sample covariance. Defaults to the population covariance.
name Python str name prefixed to Ops created by this function. Default value: None (i.e., 'covariance_reducer').

ddof

event_ndims

name

parameters

transform_fn

Methods

finalize

View source

finalize(
    final_reducer_state
)

Finalizes covariance calculation from the final_reducer_state.

Args
final_reducer_state CovarianceReducerStates that represent the final running covariance state.

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

initialize

View source

initialize(
    initial_chain_state, initial_kernel_results=None
)

Initializes a CovarianceReducerState using previously defined metadata.

For calculation purposes, the initial_chain_state does not count as a sample. This is a deliberate decision that ensures consistency across sampling procedures (i.e. tfp.mcmc.sample_chain follows similar semantics).

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). It is used to infer the shape and dtype of future samples.
initial_kernel_results A (possibly nested) structure of Tensors representing internal calculations made in a related TransitionKernel.

Returns
state CovarianceReducerState with cov_state field 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 similar to those of batching and 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 CovarianceReducerStates representing the current state of the running covariance.
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 CovarianceReducerState with updated running statistics. Its cov_state field has an identical structure to the results of self.transform_fn. Each of the individual values in that structure subsequently mimics the structure of current_reducer_state.