Maximum depth of the tree implicitly built by NUTS. The
maximum number of leapfrog steps is bounded by 2**max_tree_depth i.e.
the number of nodes in a binary tree max_tree_depth nodes deep. The
default setting of 10 takes up to 1024 leapfrog steps.
Scalar threshold of energy differences at each leapfrog,
divergence samples are defined as leapfrog steps that exceed this
threshold. Default to 1000.
The number of leapfrogs to unroll per tree
expansion step. Applies a direct linear multipler to the maximum
trajectory length implied by max_tree_depth. Defaults to 1.
The number of iterations allowed to run in parallel.
It must be a positive integer. See tf.while_loop for more details.
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.