tfp.experimental.mcmc.ProgressBarReducer

Reducer that displays a progress bar.

Inherits From: Reducer

Note this is not XLA-compatible (tf.function(experimental_compile=True)). Numpy and JAX substrates are not supported.

Example usage:

kernel = ...
current_state = ...
num_results = ...
pbar = tfp.experimental.mcmc.ProgressBarReducer(num_results)
_, final_state, kernel_results = tfp.experimental.mcmc.sample_fold(
    num_steps=num_results,
    current_state=current_state,
    kernel=kernel,
    reducer=pbar,
)

num_results Integer number of results to expect (as passed to sample chain).
progress_bar_fn A function that takes an integer num_results and returns an iterator that advances a progress bar. Defaults to tqdm progress bars (make sure they are pip installed befure using.)

num_results

parameters

progress_bar_fn

Methods

finalize

View source

Finalizes target statistic calculation from the final_state.

This is an identity function of the final_state by default. Subclasses can override it for streaming calculations whose running state is not the same as the desired result.

Args
final_reducer_state A tuple, namedtuple or list of Tensors representing the final state of the reduced statistic.

Returns
statistic An estimate of the target statistic

initialize

View source

Initialize progress bars.

All arguments are ignored.

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

Returns
state empty list.

one_step

View source

Advance progress bar by one result.

All arguments are ignored.

Args
new_chain_state A (possibly nested) structure of incoming chain state(s) with shape and dtype compatible with those used to initialize the TracingState.
current_reducer_state TracingStates representing all previously traced results.
previous_kernel_results A (possibly nested) structure of Tensors representing internal calculations made in a related TransitionKernel.

Returns
new_reducer_state empty list.