# 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
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for schedules.py."""



class LearningRateTest(absltest.TestCase):
  """Tests learning rate decay schedules."""

  def test_polynomial_decay(self):
    """Test polynomial schedule works correctly with decay_steps_factor."""
    hps = config_dict.ConfigDict(dict(
        lr_hparams={
            'schedule': 'polynomial',
            'power': 2.0,
            'base_lr': .1,
            'end_factor': .01,
            'decay_steps_factor': 0.5,
        }
    ))
    max_training_steps = 400

# 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

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for schedules.py."""



class LearningRateTest(absltest.TestCase):
  """Tests learning rate decay schedules."""

  def test_polynomial_decay(self):
    """Test polynomial schedule works correctly with decay_steps_factor."""
    hps = config_dict.ConfigDict(dict(
        lr_hparams={
            'schedule': 'polynomial',
            'power': 2.0,
            'base_lr': .1,
            'end_factor': .01,
            'decay_steps_factor': 0.5,
        }
    ))
    max_training_steps = 400

    lr_fn = schedules.get_schedule_fn(hps.lr_hparams, max_training_steps)