Skip to content

← Optax step 6 of 25

Easy Primitives

Constant Schedule

Jump to the task

A constant schedule returns the same learning rate at every step. The schedule is a function; calling it produces the number to use.

import optax

schedule = optax.constant_schedule(0.001)
first = schedule(0)       # 0.001
later = schedule(1000)    # also 0.001

An optimizer can accept that function as its learning rate, for example optax.adam(learning_rate=schedule). A fixed scalar also works when you do not need the schedule interface.

Use optax.join_schedules to join schedule functions into successive phases. optax.chain combines gradient transformations, not schedule functions. They serve different purposes.

Inputs

Implement constant_lr_at(step, lr_value) using optax.constant_schedule. Create the schedule, evaluate it and return its value.

  • step: a nonnegative whole-number step count, supplied as a scalar. Convert it to a Python integer before evaluating the schedule.
  • lr_value: the scalar learning rate to return at every step.

Output

Return lr_value as a float32 JAX scalar. Return the evaluated value, not the schedule function itself.

The chart starts with a constant rate of 0.001. Switch curves to compare how other schedules vary with the step count.

Loading visualization…