KernelSequence.end_warmup()#
- KernelSequence.end_warmup(prng_key, kernel_states, model_state, tuning_history)[source]#
Asks the kernel to inspect the warmup history and react to it.
This method is executed once the first non-warmup epoch is encountered and before
start_epoch()
is called.tuning_history
isNone
if no tuning has happened prior to the first non-warmup epoch. Otherwisetuning_history
has the same structure as returned fromtune()
but each leaf has an additional dimension. The first dimension refers to the n-the tuning.Must be jittable.
The signature is likely to change.
- Return type: