Kernel.end_warmup()

abstract Kernel.end_warmup(prng_key, kernel_state, model_state, tuning_history)

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 is None if no tuning has happened prior to the first non-warmup epoch. Otherwise tuning_history has the same structure as returned from tune() 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

WarmupOutcome[TypeVar(TKernelState, bound= Any)]