KernelSequence.start_epoch()#

KernelSequence.start_epoch(prng_key, kernel_states, model_state, epoch)[source]#

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

Return type:

list[Any]