tfp.experimental.mcmc.ProgressBarReducer

Reducer that displays a progress bar.

Inherits From: Reducer

tfp.experimental.mcmc.ProgressBarReducer(
    num_results, progress_bar_fn=make_tqdm_progress_bar_fn()
)

Used in the notebooks

Used in the tutorials

Note this is not XLA-compatible (tf.function(jit_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

finalize(
    final_reducer_state
)

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(
    initial_chain_state, initial_kernel_results=None
)

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

one_step(
    new_chain_state, current_reducer_state, previous_kernel_results
)

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.