KernelSequence.end_epoch()# KernelSequence.end_epoch(prng_key, kernel_states, model_state, epoch)[source]# Called at the end of an epoch. Must be jittable. Return type: list[Any]
