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