class_docstr: Tests different learning rate schedules .
class_docstr_tok: ["Tests","different","learning","rate","schedules","."]
docstr: Test cosine schedule works correctly.
docstr_tok: ["Test","cosine","schedule","works","correctly","."]
code: def test_cosine_decay(self):
      """Test cosine schedule works correctly."""
      config = ml_collections.ConfigDict(
          dict(
              lr_configs={
                  "learning_rate_schedule": "compound",
                  "factors": "cosine_decay",
                  "steps_per_cycle": 100,
                  "t_mul": 2.0,
                  "m_mul": 0.5,
                  "alpha": 0.3,
                  "base_learning_rate": 1.0,
              }
          )
      )
      lr_fn = lr_schedules.get_learning_rate_fn(config)
      config = config.lr_configs
      tf_cosine_decay = tf.keras.experimental.CosineDecayRestarts(
          initial_learning_rate=config["base_learning_rate"],
          first_decay_steps=config["steps_per_cycle"],
          t_mul=config["t_mul"],
          m_mul=config["m_mul"],
          alpha=config["alpha"],
      )
      for step in range(400):
          expected_learning_rate = tf_cosine_decay(step=step).numpy()
          self.assertAlmostEqual(lr_fn(step), expected_learning_rate)
  
file_path: scenic/scenic/train_lib/tests/test_lr_schedules.py
imports: from absl.testing import absltest
  import ml_collections
  from scenic.train_lib import lr_schedules
  import tensorflow
label_desc: Use image node_id 5 for calling the LearningRateScchedulesTest obj's underlying member method code with example usage: obj.test_cosine_decay() without return types