This uses Hamiltonian Monte Carlo to do the sampling. Step size is tuned using
a dual-averaging adaptation, and the kernel is conditioned using a diagonal
mass matrix, which is estimated using expanding windows.
Number of draws after adaptation.
A joint distribution to sample from.
Number of leapfrog steps to use for the Hamiltonian Monte Carlo step.
Number of independent chains to run MCMC with.
The trace function should accept the arguments
(state, bijector, is_adapting, phmc_kernel_results), where the state
is an unconstrained, flattened float tensor, bijector is the
tfb.Bijector that is used for unconstraining and flattening,
is_adapting is a boolean to mark whether the draw is from an adaptation
step, and phmc_kernel_results is the
UncalibratedPreconditionedHamiltonianMonteCarloKernelResults from the
PreconditionedHamiltonianMonteCarlo kernel. Note that
bijector.inverse(state) will provide access to the current draw in the
untransformed space, using the structure of the provided joint_dist.
If True, then the final kernel results are
returned alongside the chain state and the trace specified by the
Whether to return tuning traces and draws.
Optional, a seed for reproducible sampling.
These are used to condition the provided joint distribution, and are
passed directly to joint_dist.experimental_pin(**pins).
A single structure of draws is returned in case the trace_fn is None, and
return_final_kernel_results is False. If there is a trace function,
the return value is a tuple, with the trace second. If the
return_final_kernel_results is True, the return value is a tuple of
length 3, with final kernel results returned last. If discard_tuning is
True, the tensors in draws and trace will have length n_draws,
otherwise, they will have length n_draws + num_adaptation_steps.