tfp.experimental.mcmc.WithReductions

Applies Reducers to stream over MCMC samples.

Inherits From: TransitionKernel

Used in the notebooks

Used in the tutorials

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 reduction_results 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.
adjust_kr_fn Optional function to adjust the kernel_results structure of inner_kernel before presenting it to reducer. Useful for drivers (like sample_fold) that construct their own kernel onions, but accept Reducers as arguments.
name Python str name prefixed to Ops created by this function. Default value: None (i.e., "reduced_kernel").

adjust_kr_fn

experimental_shard_axis_names The shard axis names for members of the state.
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 reduction_results 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).

experimental_with_shard_axes

View source

Returns a copy of the kernel with the provided shard axis names.

Args
shard_axis_names a structure of strings indicating the shard axis names for each component of this kernel's state.

Returns
A copy of the current kernel with the shard axis information.

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 reduction_results and a reference to kernel results of nested TransitionKernels.
seed PRNG seed; see tfp.random.sanitize_seed for details.

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