# Path: vmoe/train/schedule.py
# def create_learning_rate_schedule(*,
#                                   schedule: str,
#                                   total_steps: int,
#                                   **kwargs) -> optax.Schedule:
# def warmup_polynomial_decay_schedule(peak_value: float, end_value: float,
#                                      power: float, warmup_steps: int,
#                                      decay_steps: int) -> optax.Schedule:


from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import schedule
import chex
import numpy as np
        **extra_kwargs,
    }
    fn = schedule.create_learning_rate_schedule(**kwargs, total_steps=8)
    values = np.asarray(list(map(fn, range(10))), dtype=np.float32)
    self.assertAlmostEqual(values[0], 0., places=5)
    self.assertAlmostEqual(values[4], 0.1, places=5)
    self.assertAlmostEqual(values[8], 0.01, places=5)
    self.assertAlmostEqual(values[9], 0.01, places=5)

  def test_create_learning_rate_warmup_exponential_decay(self):
    kwargs = {
        'schedule': 'warmup_exponential_decay',
        'peak_value': 0.1,
        'warmup_steps': 4,
        'transition_steps': 1,
        'decay_rate': 0.9,
    }
    fn = schedule.create_learning_rate_schedule(**kwargs, total_steps=8)
    values = np.asarray(list(map(fn, range(10))), dtype=np.float32)
    self.assertAlmostEqual(values[0], 0., places=5)
    self.assertAlmostEqual(values[4], 0.1, places=5)
    self.assertAlmostEqual(values[8], 0.1 * 0.9**4, places=5)
    self.assertAlmostEqual(values[9], 0.1 * 0.9**5, places=5)

  def test_create_learning_rate_schedule_raises(self):
    kwargs = {'schedule': 'foo_bar'}
    with self.assertRaisesRegex(ValueError, 'Unknown learning rate schedule'):
      schedule.create_learning_rate_schedule(**kwargs, total_steps=10)

  def test_warmup_polynomial_decay_schedule(self):

# Path: vmoe/train/schedule.py
# def create_learning_rate_schedule(*,
#                                   schedule: str,
#                                   total_steps: int,
#                                   **kwargs) -> optax.Schedule:
# def warmup_polynomial_decay_schedule(peak_value: float, end_value: float,
#                                      power: float, warmup_steps: int,
#                                      decay_steps: int) -> optax.Schedule:


# Path: vmoe/train/schedule_test.py
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import schedule
import chex
import numpy as np

        **extra_kwargs,
    }
    fn = schedule.create_learning_rate_schedule(**kwargs, total_steps=8)
    values = np.asarray(list(map(fn, range(10))), dtype=np.float32)
    self.assertAlmostEqual(values[0], 0., places=5)
    self.assertAlmostEqual(values[4], 0.1, places=5)
    self.assertAlmostEqual(values[8], 0.01, places=5)
    self.assertAlmostEqual(values[9], 0.01, places=5)

  def test_create_learning_rate_warmup_exponential_decay(self):
    kwargs = {
        'schedule': 'warmup_exponential_decay',
        'peak_value': 0.1,
        'warmup_steps': 4,
        'transition_steps': 1,
        'decay_rate': 0.9,
    }
    fn = schedule.create_learning_rate_schedule(**kwargs, total_steps=8)
    values = np.asarray(list(map(fn, range(10))), dtype=np.float32)
    self.assertAlmostEqual(values[0], 0., places=5)
    self.assertAlmostEqual(values[4], 0.1, places=5)
    self.assertAlmostEqual(values[8], 0.1 * 0.9**4, places=5)
    self.assertAlmostEqual(values[9], 0.1 * 0.9**5, places=5)

  def test_create_learning_rate_schedule_raises(self):
    kwargs = {'schedule': 'foo_bar'}
    with self.assertRaisesRegex(ValueError, 'Unknown learning rate schedule'):
      schedule.create_learning_rate_schedule(**kwargs, total_steps=10)

  def test_warmup_polynomial_decay_schedule(self):

    fn = schedule.warmup_polynomial_decay_schedule(