TransitionMixin.transition()
TransitionMixin.transition()#
- TransitionMixin.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
)]