# Path: init2winit/schedules.py
# def _check_schedule_hparams(schedule_hparams, expected_keys):
# def constant_schedule(schedule_hparams, max_training_steps):
# def cosine_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def rsqrt_normalized_decay(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def polynomial_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def piecewise_constant_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def piecewise_linear_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def mlperf_polynomial_schedule(schedule_hparams, max_training_steps):
#   def step_fn(step):
# def compound_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(step):
# def prepend_linear_warmup(schedule_hparams, max_training_steps,
#                           base_lr_schedule):
#   def lr_fn(t):
# def warmup_then_piecewise_constant_schedule(schedule_hparams,
#                                             max_training_steps):
# def schedule_stretcher(schedule_fn, stretch_factor):
#   def _schedule(global_step):
# def get_schedule_fn(schedule_hparams, max_training_steps, stretch_factor=1):


from absl.testing import absltest
from init2winit import schedules
from ml_collections.config_dict import config_dict
import tensorflow.compat.v1 as tf
        }
    ))
    max_training_steps = 400
    lr_fn = schedules.get_schedule_fn(hps.lr_hparams, max_training_steps)
    hps = hps.lr_hparams
    decay_steps = hps['decay_steps']
    for step in range(max_training_steps):
      expected_learning_rate = tf.train.polynomial_decay(
          hps['base_lr'],
          step,
          decay_steps,
          hps['end_factor'] * hps['base_lr'],
          power=hps['power'])().numpy()
      self.assertAlmostEqual(lr_fn(step), expected_learning_rate)

  def test_schedule_stretching(self):
    """Test that schedules can be properly stretched."""
    max_training_steps = 100
    lr_hparams = config_dict.ConfigDict({
        'schedule': 'mlperf_polynomial',
        'base_lr': 10.0,
        'warmup_steps': 10,
        'decay_end': -1,
        'end_lr': 1e-4,
        'power': 2.0,
        'start_lr': 0.0,
        'warmup_power': 1.0,
    })
    lr_fn = schedules.get_schedule_fn(lr_hparams, max_training_steps)
    stretch_factor = 3

# Path: init2winit/schedules.py
# def _check_schedule_hparams(schedule_hparams, expected_keys):
# def constant_schedule(schedule_hparams, max_training_steps):
# def cosine_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def rsqrt_normalized_decay(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def polynomial_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def piecewise_constant_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def piecewise_linear_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(t):
# def mlperf_polynomial_schedule(schedule_hparams, max_training_steps):
#   def step_fn(step):
# def compound_schedule(schedule_hparams, max_training_steps):
#   def lr_fn(step):
# def prepend_linear_warmup(schedule_hparams, max_training_steps,
#                           base_lr_schedule):
#   def lr_fn(t):
# def warmup_then_piecewise_constant_schedule(schedule_hparams,
#                                             max_training_steps):
# def schedule_stretcher(schedule_fn, stretch_factor):
#   def _schedule(global_step):
# def get_schedule_fn(schedule_hparams, max_training_steps, stretch_factor=1):


# Path: init2winit/test_schedules.py
from absl.testing import absltest
from init2winit import schedules
from ml_collections.config_dict import config_dict
import tensorflow.compat.v1 as tf

        }
    ))
    max_training_steps = 400
    lr_fn = schedules.get_schedule_fn(hps.lr_hparams, max_training_steps)
    hps = hps.lr_hparams
    decay_steps = hps['decay_steps']
    for step in range(max_training_steps):
      expected_learning_rate = tf.train.polynomial_decay(
          hps['base_lr'],
          step,
          decay_steps,
          hps['end_factor'] * hps['base_lr'],
          power=hps['power'])().numpy()
      self.assertAlmostEqual(lr_fn(step), expected_learning_rate)

  def test_schedule_stretching(self):
    """Test that schedules can be properly stretched."""
    max_training_steps = 100
    lr_hparams = config_dict.ConfigDict({
        'schedule': 'mlperf_polynomial',
        'base_lr': 10.0,
        'warmup_steps': 10,
        'decay_end': -1,
        'end_lr': 1e-4,
        'power': 2.0,
        'start_lr': 0.0,
        'warmup_power': 1.0,
    })
    lr_fn = schedules.get_schedule_fn(lr_hparams, max_training_steps)
    stretch_factor = 3

    stretched_lr_fn = schedules.get_schedule_fn(