Fix the learning rate warmup schedule

Review the AI's codeEasy10 minOpenAI and Anthropic style

Your task

An AI wrote a training loop with learning rate warmup, where the rate climbs from near zero to base_lr over the first steps and then holds. Its loop started at step 1, stopped one step early, and the rate never reached its peak. Your team writes the schedule out as a list, one rate per step, so it can be checked.

An assistant wrote warmup_lr_schedule(total_steps, warmup_steps, base_lr). It is in the editor. It has a bug. Find it and fix it.

Input and output:

  • total_steps: how many training steps run.
  • warmup_steps: how many steps the warmup takes.
  • base_lr: the peak learning rate.
  • Return the learning rate of each step, in order.

Rules the tests check:

  • Return exactly total_steps rates, or an empty list if total_steps is 0 or less.
  • Number the steps from 1. Step k's rate is base_lr times the smaller of k divided by warmup_steps, and 1. So the first step already trains, and the rate reaches base_lr exactly at step warmup_steps.
  • If warmup_steps is 0 or less, every rate is base_lr.
  • Answers within 0.000001 of the expected values pass.

Examples

warmup_lr_schedule(total_steps: int, warmup_steps: int, base_lr: float) -> list[float]
Example 1
Inputwarmup_lr_schedule(6, 4, 1)
Output[0.25, 0.5, 0.75, 1, 1, 1]
Example 2
Inputwarmup_lr_schedule(3, 5, 0.5)
Output[0.1, 0.2, 0.3]