code: def test_schedulers(self):
      common_kwargs = {"num_warmup_steps": 2, "num_training_steps": 10}
      # schedulers doct format
      # function: (sched_args_dict, expected_learning_rates)
      scheds = {
          get_constant_schedule: ({}, [10.0] * self.num_steps),
          get_constant_schedule_with_warmup: (
              {"num_warmup_steps": 4},
              [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
          ),
          get_linear_schedule_with_warmup: (
              {**common_kwargs},
              [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
          ),
          get_cosine_schedule_with_warmup: (
              {**common_kwargs},
              [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
          ),
          get_cosine_with_hard_restarts_schedule_with_warmup: (
              {**common_kwargs, "num_cycles": 2},
              [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
          ),
          get_polynomial_decay_schedule_with_warmup: (
              {**common_kwargs, "power": 2.0, "lr_end": 1e-7},
              [
                  0.0,
                  5.0,
                  10.0,
                  7.656,
                  5.625,
                  3.906,
                  2.5,
                  1.406,
                  0.625,
                  0.156,
              ],
          ),
      }
  
      for scheduler_func, data in scheds.items():
          kwargs, expected_learning_rates = data
  
          scheduler = scheduler_func(self.optimizer, **kwargs)
          self.assertEqual(len([scheduler.get_lr()[0]]), 1)
          lrs_1 = unwrap_schedule(scheduler, self.num_steps)
          self.assertListAlmostEqual(
              lrs_1,
              expected_learning_rates,
              tol=1e-2,
              msg=f"failed for {scheduler_func} in normal scheduler",
          )
  
          scheduler = scheduler_func(self.optimizer, **kwargs)
          lrs_2 = unwrap_and_save_reload_schedule(
              scheduler, self.num_steps
          )
          self.assertListEqual(
              lrs_1,
              lrs_2,
              msg=f"failed for {scheduler_func} in save and reload",
          )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/tests/optimization/test_optimization.py
imports: import os
  import tempfile
  import unittest
  from transformers import is_torch_available
  from transformers.testing_utils import require_torch
label_desc: Use image node_id 2 for calling the ScheduleInitTest obj's underlying member method code with example usage: obj.test_schedulers() without return types