Skip to content

← Optax step 9 of 25

Medium Primitives

Piecewise Constant Schedule

Jump to the task

A piecewise constant schedule holds the learning rate steady between chosen steps. At each boundary, it multiplies the current rate by a scale factor. The factors accumulate as you pass more boundaries.

The fixed schedule for this exercise

Start at 0.1, multiply by 0.5 at step 100, then multiply by 0.5 again at step 200:

Step range Learning rate
0 through 99 0.1
100 through 199 0.1 * 0.5 = 0.05
200 onward 0.1 * 0.5 * 0.5 = 0.025

The change takes effect at the boundary itself. Step 99 uses 0.1; step 100 already uses 0.05. Step 199 uses 0.05, and step 200 uses 0.025. There are no further changes after the second boundary.

Multipliers, not target values

Create optax.piecewise_constant_schedule with init_value=0.1 and boundaries_and_scales={100: 0.5, 200: 0.5}. The dictionary maps each boundary step to its multiplier. A value of 0.5 means "halve the rate", not "set the rate to 0.5".

Evaluating the schedule does not advance a counter or change its state. Each call calculates the value for the step you supply. Calling it at step 200 does not require calling it at steps 0 through 199 first.

Inputs

Implement piecewise_lr_at(step) using the fixed Optax schedule above.

  • step: a nonnegative whole-number step count, supplied as a scalar. Convert it to a Python integer before evaluating the schedule.

Output

Return the learning rate as a float32 JAX scalar. The initial value, boundaries and multipliers are fixed inside the function.

The chart below compares schedule shapes using separate demonstration settings. Its step curve starts at 0.001 and drops by a factor of ten at each third of the displayed range. It does not plot this exercise's two halving boundaries.

Loading visualization…