tfp.experimental.mcmc.ProgressBarReducer

See Stable See Nightly

View source on GitHub

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
TFP Release Notes notebook (0.12.1)

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,
)

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

Attributes
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.