tfp.experimental.mcmc.WithReductions

Applies Reducers to stream over MCMC samples.

Inherits From: TransitionKernel

tfp.experimental.mcmc.WithReductions(
    inner_kernel, reducer, adjust_kr_fn=(lambda x: x), name=None
)

Used in the notebooks

Used in the tutorials
TFP Release Notes notebook (0.12.1)

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.

Args
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").

Attributes
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

bootstrap_results(
    init_state, inner_results=None, previous_reducer_state=None
)

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

copy(
    **override_parameter_kwargs
)

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

experimental_with_shard_axes(
    shard_axis_names
)

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

one_step(
    current_state, previous_kernel_results, seed=None
)

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.