tfp.experimental.mcmc.WithReductions

Applies Reducers to stream over MCMC samples.

Inherits From: TransitionKernel

WithReductions augments an inner MCMC kernel with side-computations that can read the stream of samples as they are generated. This is relevant for streaming uses of MCMC, where materializing the entire Markov chain history is undesirable, e.g. due to memory limits.

One WithReductions instance can attach an arbitrary collection of side-computations, each of which must be packaged as a Reducer. WithReductions operates by generating a sample with its inner_kernel's one_step, then invoking each Reducer's one_step method on that sample. The updated reducer states are stored in the streaming_calculations field of WithReductions' kernel results.

inner_kernel TransitionKernel whose one_step will generate MCMC sample(s).
reducer A (possibly nested) structure of Reducers to be evaluated on the inner_kernel's samples.
name Python str name prefixed to Ops created by this function. Default value: None (i.e., "reduced_kernel").

inner_kernel

is_calibrated Returns True if Markov chain converges to specified distribution.

TransitionKernels which are "uncalibrated" are often calibrated by composing them with the tfp.mcmc.MetropolisHastings TransitionKernel.

name

parameters

reducer

Methods

bootstrap_results

View source

Instantiates reducer states with identical structure to the init_state.

Args
init_state Tensor or Python list of Tensors representing the state(s) of the Markov chain(s). For consistency across sampling procedures (i.e. tfp.mcmc.sample_chain follows similar semantics), the initial state does not count as a "sample". Hence, all reducer states will reflect empty streams.
inner_results Optional results tuple for the inner kernel. Will be re-bootstrapped if omitted.
previous_reducer_state Optional results structure for the reducers. Will be re-initialized if omitted.

Returns
kernel_results WithReductionsKernelResults representing updated kernel results. Reducer states are stored in the streaming_calculations field. The state structure is identical to self.reducer.

copy

View source

Non-destructively creates a deep copy of the kernel.

Args
**override_parameter_kwargs Python String/value dictionary of initialization arguments to override with new values.

Returns
new_kernel TransitionKernel object of same type as self, initialized with the union of self.parameters and override_parameter_kwargs, with any shared keys overridden by the value of override_parameter_kwargs, i.e., dict(self.parameters, **override_parameters_kwargs).

one_step

View source

Updates all Reducers with a new sample from the inner_kernel.

Args
current_state Tensor or Python list of Tensors representing the current state(s) of the Markov chain(s),
previous_kernel_results WithReductionsKernelResults named tuple. WithReductionsKernelResults contain the state of streaming_calculations and a reference to kernel results of nested TransitionKernels.
seed Optional seed for reproducible sampling.

Returns
new_state Newest MCMC state drawn from the inner_kernel.
kernel_results WithReductionsKernelResults representing updated kernel results. Reducer states are stored in the streaming_calculations field. The state structure is identical to self.reducer.