tfp.experimental.mcmc.default_nuts_trace_fn

Trace function for windowed_adaptive_nuts providing standard diagnostics.

tfp.experimental.mcmc.default_nuts_trace_fn(
    state, bijector, is_adapting, pkr
)

Specifically, these match up with a number of diagnostics used by ArviZ [1], to make diagnostics and analysis easier. The names used follow those used in TensorFlow Probability, and will need to be mapped to those used in the ArviZ schema.

[1]: Kumar, R., Carroll, C., Hartikainen, A., & Martin, O. (2019). ArviZ a unified library for exploratory analysis of Bayesian models in Python. Journal of Open Source Software, 4(33), 1143.

state tf.Tensor Current sampler state, flattened and unconstrained.
bijector tfb.Bijector This can be used to untransform the shape to something with the same shape as will be returned.
is_adapting bool Whether this is an adapting step, or may be treated as a valid MCMC draw.
pkr UncalibratedPreconditionedHamiltonianMonteCarloKernelResults Kernel results from this iteration.

dict with sampler statistics.