NUTSKernel#
- class liesel.goose.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, identifier='')[source]#
Bases:
ModelMixin,TransitionMixin[NUTSKernelState,NUTSTransitionInfo],TuningMixin[NUTSKernelState,DefaultTuningInfo],ReprMixinA NUTS kernel with dual averaging and an inverse mass matrix tuner, implementing the
Kernelprotocol.Methods
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.
init_state(prng_key, model_state)Initializes the kernel state with an identity inverse mass matrix and a reasonable step size (unless explicit arguments were provided by the user).
start_epoch(prng_key, kernel_state, ...)Resets the state of the dual averaging algorithm.
Attributes
Dict of error codes and their meaning.
Kernel identifier, set by
EngineBuilderWhether this kernel needs its history for tuning.
Tuple of position keys handled by this kernel.