Study Plan

Traced Learning Rate Schedule

JAX

Medium

Compute a learning rate that rises linearly during warmup and then follows a cosine curve down to a floor, for a step count that is a traced JAX array.

For an entry $s$ of step, with $W$ equal to warmup_steps, $S$ equal to total_steps, $\eta_{\max}$ equal to peak_lr and $\eta_{\min}$ equal to min_lr, the learning rate is

\eta(s) = \begin{cases} \eta_{\max} \, \dfrac{s}{W} & s < W \\[2mm] \eta_{\min} + \dfrac{\eta_{\max} - \eta_{\min}}{2} \left(1 + \cos(\pi p)\right) & s \ge W \end{cases}, \qquad p = \min\left(\frac{s - W}{S - W},\ 1\right).

Your function is compiled with jax.jit, with step traced and the other four arguments static, so step is a placeholder without a known value while the code runs. A Python if statement, float(step) or int(step) cannot be used on it.

Return a float32 JAX array with the same shape as step, holding the learning rate of each step.

Examples

Input: step = [0,5,10,55,100,150], peak_lr = 1.0, min_lr = 0.0, warmup_steps = 10, total_steps = 100

Output: [0.0,0.5,1.0,0.5,0.0,0.0]

Explanation: Warmup covers steps 0 to 9, so step 5 is halfway up the ramp at 0.5 and step 10 reaches the peak of 1. Step 55 is halfway through the decay, where the cosine is 0 and the rate is the midpoint 0.5, and steps 100 and 150 sit at the floor of 0.

Input: step = 3, peak_lr = 0.001, min_lr = 0.0001, warmup_steps = 8, total_steps = 40

Output: 0.000375

Input: step = [0,1,2,3], peak_lr = 0.5, min_lr = 0.05, warmup_steps = 1, total_steps = 3

Output: [0.0,0.5,0.275,0.05]

Hint 1

jnp.where(step < warmup, a, b) replaces a Python if on a traced step.

Hint 2

jnp.clip(progress, 0.0, 1.0) keeps the decay progress in range.

Hint 3

step.astype(jnp.float32) before dividing.

Requirements

Choose the phase for every entry of step separately with jnp.where or an equivalent selection, because step may hold one step or several and jax.lax.cond accepts only a single scalar condition.

Constraints

step is an int32 JAX array of shape () or (K,) with K from 1 to 6 and values from 0 to 4500.
peak_lr and min_lr are Python floats with 0 <= min_lr <= peak_lr <= 2.
warmup_steps and total_steps are Python ints with 1 <= warmup_steps < total_steps <= 3000.

Try Similar Problems

Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Static Arguments in jitEasy Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium Global Norm Gradient ClippingMedium AdamW from ScratchMedium Mixed Precision with Loss ScalingMedium