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)]