Kernel.transition()
Kernel.transition()#
- abstract Kernel.transition(prng_key, kernel_state, model_state, epoch)[source]#
Handles one transition. Must be jittable.
- Return type
TransitionOutcome
[TypeVar
(TKernelState
, bound=Any
),TypeVar
(TTransitionInfo
, bound=TransitionInfo
)]