tfp.substrates.jax.mcmc.StatesAndTrace

States and auxiliary trace of an MCMC chain.

The first dimension of all the Tensors in this structure is the same and represents the chain length.

all_states A Tensor or a nested collection of Tensors representing the MCMC chain state.
trace A Tensor or a nested collection of Tensors representing the auxiliary values traced alongside the chain.