- class liesel.goose.nuts.NUTSKernel(position_keys, initial_step_size=None, initial_inverse_mass_matrix=None, max_treedepth=10, da_target_accept=0.8, da_gamma=0.05, da_kappa=0.75, da_t0=10, mm_diag=True)#
A NUTS kernel with dual averaging and an inverse mass matrix tuner, implementing the
end_epoch(prng_key, kernel_state, ...)
Sets the step size as found by the dual averaging algorithm.
end_warmup(prng_key, kernel_state, ...)
Currently does nothing.
Whether the model interface is set.
Initializes the kernel state with an identity inverse mass matrix and a reasonable step size (unless explicit arguments were provided by the user).
Returns the log-probability function with the position as the only argument.
Extracts the position from a model state.
Sets the model interface.
start_epoch(prng_key, kernel_state, ...)
Resets the state of the dual averaging algorithm.
transition(prng_key, kernel_state, ...)
Handles one transition.
tune(prng_key, kernel_state, model_state, ...)
The method can perform automatic tuning of the kernel and is called after each adaptation epoch.
Dict of error codes and their meaning.
Kernel identifier, set by
Returns the model interface if it is set.
Whether this kernel needs its history for tuning.
Tuple of position keys handled by this kernel.