def test_key_values_correct(
      self,
      step: int,
      value: float,
      learning_rate: float,
      warmup: bool,
      warmup_steps: int,
      linear_decay: bool,
      max_steps: int,
      decay_minimum_factor: float,
  ):
    """Match scheduler output with expected value."""
    lr_scheduler = optim_utils.create_learning_rate_scheduler(
        learning_rate=learning_rate,
        warmup=warmup,
        warmup_steps=warmup_steps,
        linear_decay=linear_decay,
        max_steps=max_steps,
        decay_minimum_factor=decay_minimum_factor,
    )
    self.assertEqual(lr_scheduler(step), value)
Match scheduler output with expected value.
language/mentionmemory/utils/optim_utils_test.py