Kernel.end_epoch()#

abstract Kernel.end_epoch(prng_key, kernel_state, model_state, epoch)[source]#

Called at the end of an epoch. Must be jittable.

Return type

TypeVar(TKernelState, bound= Any)