def create_learning_rate_scheduler(
    learning_rate: float,
    warmup: bool = False,
    warmup_steps: int = 1000,
    linear_decay: bool = False,
    max_steps: int = 500000,
    decay_minimum_factor: float = 0.0,
) -> Callable[[int], float]:
  """Creates learning rate scheduler with options for warmup and linear decay.

  Args:
    learning_rate: base learning rate to be modified.
    warmup: if true, applies learning rate warmup.
    warmup_steps: number of warmup steps.
    linear_decay: if true, applies linear learning rate decay.
    max_steps: number of steps after which learning rate is fully decayed.
    decay_minimum_factor: proportion of lr left after max_steps.

  Returns:
    Function that maps from step to lr.
  """

  def step_fn(step: int) -> float:
    factor = 1.0
    if warmup:
      factor *= jnp.minimum(1.0, step / warmup_steps)
    if linear_decay:
      step_decay = (1.0 - decay_minimum_factor) / (max_steps - warmup_steps)
      factor *= jnp.minimum(1.0, 1.0 - step_decay * (step - warmup_steps))
    return jnp.asarray(learning_rate * factor)  # pytype: disable=bad-return-type  # jnp-type

  return step_fn
Creates learning rate scheduler with options for warmup and linear decay.

  Args:
    learning_rate: base learning rate to be modified.
    warmup: if true, applies learning rate warmup.
    warmup_steps: number of warmup steps.
    linear_decay: if true, applies linear learning rate decay.
    max_steps: number of steps after which learning rate is fully decayed.
    decay_minimum_factor: proportion of lr left after max_steps.

  Returns:
    Function that maps from step to lr.
  
language/mentionmemory/utils/optim_utils.py