API Reference
Called at the beginning of an epoch. Must be jittable.
TypeVar(TKernelState, bound= Any)
TypeVar
TKernelState
Any