tests for learning rate schedule api import math import numpy as np import pytest from keras import backend from keras import layers from keras import optimizers from keras import testing from keras models import sequential from keras optimizers import schedules class testfitlrschedulesflowtesting testcase pytest mark requirestrainablebackend def testfitlrcorrectnessself model sequential layers dense 2 kernelinitializerones biasinitializerones optimizer optimizers adam learningrateschedules exponentialdecay initiallearningrate0 05 decaysteps1 decayrate0 9 self assertequallenoptimizer variables 1 self assertequaloptimizer variables0 0 model compileoptimizeroptimizer lossmse x np arange32 reshape16 2 y np arange32 reshape16 2 history model fitx y epochs3 batchsize4 shufflefalse self assertequaloptimizer variables0 4 3 self assertallclose history historyloss 230 79457092285156 128 30319213867188 79 33648681640625 rtol5e5 class exponentialdecaytesttesting testcase def testconfigself self runclassserializationtest schedules exponentialdecay initiallearningrate0 05 decaysteps10 decayrate0 96 staircasetrue namemyed def testcontinuousself step 5 decayedlr schedules exponentialdecay0 05 10 0 96 expected 0 05 0 96 5 0 10 0 self assertallclosedecayedlrstep expected 1e6 def teststaircaseself step backend variable1 0 decayedlr schedules exponentialdecay0 1 3 0 96 staircasetrue no change to learning rate due to staircase expected 0 1 self assertallclosedecayedlrstep expected 1e6 expected 0 1 step assign2 self assertallclosedecayedlrstep expected 1e6 decayed learning rate expected 0 1 0 96 100 3 step assign100 self assertallclosedecayedlrstep expected 1e6 def testvariablesself step backend variable1 0 decayedlr schedules exponentialdecay0 1 3 0 96 staircasetrue no change to learning rate step assign1 self assertallclosedecayedlrstep 0 1 1e6 step assign2 self assertallclosedecayedlrstep 0 1 1e6 decayed learning rate step assign100 expected 0 1 0 96 100 3 self assertallclosedecayedlrstep expected 1e6 class piecewiseconstantdecaytesttesting testcase def testconfigself self runclassserializationtest schedules piecewiseconstantdecay boundaries10 20 values1 2 3 namemypcd def testpiecewisevaluesself x backend variable999 0 decayedlr schedules piecewiseconstantdecay 100 110 120 1 0 0 1 0 01 0 001 self assertallclosedecayedlrx 1 0 1e6 x assign100 self assertallclosedecayedlrx 1 0 1e6 x assign105 self assertallclosedecayedlrx 0 1 1e6 x assign110 self assertallclosedecayedlrx 0 1 1e6 x assign120 self assertallclosedecayedlrx 0 01 1e6 x assign999 self assertallclosedecayedlrx 0 001 1e6 def testboundaryvaluesself test casting boundaries from int32 to int64 xint64 backend variable0 dtypeint64 trainablefalse boundaries values 1 2 3 0 4 0 5 0 6 0 7 decayedlr schedules piecewiseconstantdecayboundaries values self assertallclosedecayedlrxint64 0 4 1e6 xint64 assign1 self assertallclosedecayedlrxint64 0 4 1e6 xint64 assign2 self assertallclosedecayedlrxint64 0 5 1e6 xint64 assign3 self assertallclosedecayedlrxint64 0 6 1e6 xint64 assign4 self assertallclosedecayedlrxint64 0 7 1e6 class lineardecaytesttesting testcase def testconfigself self runclassserializationtest schedules polynomialdecay initiallearningrate0 1 decaysteps100 endlearningrate0 005 power1 0 cyclefalse namemyld def testhalfwayself step 5 lr 0 05 endlr 0 0 decayedlr schedules polynomialdecaylr 10 endlr expected lr 0 5 self assertallclosedecayedlrstep expected 1e6 def testendself step 10 lr 0 05 endlr 0 001 decayedlr schedules polynomialdecaylr 10 endlr expected endlr self assertallclosedecayedlrstep expected 1e6 def testhalfwaywithendself step 5 lr 0 05 endlr 0 001 decayedlr schedules polynomialdecaylr 10 endlr expected lr endlr 0 5 self assertallclosedecayedlrstep expected 1e6 def testbeyondendself step 15 lr 0 05 endlr 0 001 decayedlr schedules polynomialdecaylr 10 endlr expected endlr self assertallclosedecayedlrstep expected 1e6 def testbeyondendwithcycleself step 15 lr 0 05 endlr 0 001 decayedlr schedules polynomialdecaylr 10 endlr cycletrue expected lr endlr 0 25 endlr self assertallclosedecayedlrstep expected 1e6 class sqrtdecaytesttesting testcase def testhalfwayself step 5 lr 0 05 endlr 0 0 power 0 5 decayedlr schedules polynomialdecaylr 10 endlr powerpower expected lr 0 5power self assertallclosedecayedlrstep expected 1e6 def testendself step 10 lr 0 05 endlr 0 001 power 0 5 decayedlr schedules polynomialdecaylr 10 endlr powerpower expected endlr self assertallclosedecayedlrstep expected 1e6 def testhalfwaywithendself step 5 lr 0 05 endlr 0 001 power 0 5 decayedlr schedules polynomialdecaylr 10 endlr powerpower expected lr endlr 0 5power endlr self assertallclosedecayedlrstep expected 1e6 def testbeyondendself step 15 lr 0 05 endlr 0 001 power 0 5 decayedlr schedules polynomialdecaylr 10 endlr powerpower expected endlr self assertallclosedecayedlrstep expected 1e6 def testbeyondendwithcycleself step 15 lr 0 05 endlr 0 001 power 0 5 decayedlr schedules polynomialdecay lr 10 endlr powerpower cycletrue expected lr endlr 0 25power endlr self assertallclosedecayedlrstep expected 1e6 def testbeginwithcycleself lr 0 001 decaysteps 10 step 0 decayedlr schedules polynomialdecaylr decaysteps cycletrue expected lr self assertallclosedecayedlrstep expected 1e6 class inversetimedecaytesttesting testcase def testconfigself self runclassserializationtest schedules inversetimedecay initiallearningrate0 05 decaysteps10 decayrate0 96 staircasetrue namemyitd def testdecayself initiallr 0 1 k 10 decayrate 0 96 step backend variable0 0 decayedlr schedules inversetimedecayinitiallr k decayrate for i in rangek 1 expected initiallr 1 i k decayrate self assertallclosedecayedlrstep expected 1e6 step assignstep 1 def teststaircaseself initiallr 0 1 k 10 decayrate 0 96 step backend variable0 0 decayedlr schedules inversetimedecay initiallr k decayrate staircasetrue for i in rangek 1 expected initiallr 1 decayrate i k self assertallclosedecayedlrstep expected 1e6 step assignstep 1 class cosinedecaytesttesting testcase def testconfigself self runclassserializationtest schedules cosinedecay initiallearningrate0 05 decaysteps10 alpha0 1 warmuptarget0 2 warmupsteps2 namemycd def npcosinedecayself step decaysteps alpha0 0 step minstep decaysteps completedfraction step decaysteps decay 0 5 1 0 math cosmath pi completedfraction return 1 0 alpha decay alpha def testdecayself numtrainingsteps 1000 initiallr 1 0 for step in range0 1500 250 decayedlr schedules cosinedecayinitiallr numtrainingsteps expected self npcosinedecaystep numtrainingsteps self assertallclosedecayedlrstep expected 1e6 def linearwarmupself step warmupsteps initiallr targetlr completedfraction step warmupsteps totaldelta targetlr initiallr return completedfraction totaldelta def testwarmupself warmupsteps 1500 initiallr 0 0 targetlr 10 0 for step in range0 1500 250 lr schedules cosinedecay initiallr 10 warmuptargettargetlr warmupstepswarmupsteps expected self linearwarmup step warmupsteps initiallr targetlr self assertallcloselrstep expected def testalphaself numtrainingsteps 1000 initiallr 1 0 alpha 0 1 for step in range0 1500 250 decayedlr schedules cosinedecay initiallr numtrainingsteps alpha expected self npcosinedecaystep numtrainingsteps alpha self assertallclosedecayedlrstep expected 1e6 def testfloat64self numtrainingsteps 1000 initiallr np float641 0 for step in range0 1500 250 decayedlr schedules cosinedecayinitiallr numtrainingsteps expected self npcosinedecaystep numtrainingsteps self assertallclosedecayedlrstep expected 1e6 def testwarmupdecayself warmupsteps 2000 decaysteps 1000 initiallr 0 0 targetlr 10 0 for step in range0 3000 250 lr schedules cosinedecay initiallr decaysteps warmuptargettargetlr warmupstepswarmupsteps if step warmupsteps 1 expected self linearwarmup step warmupsteps initiallr targetlr else expected targetlr self npcosinedecay step warmupsteps decaysteps self assertallcloselrstep expected class cosinedecayrestartstesttesting testcase def testconfigself self runclassserializationtest schedules cosinedecayrestarts initiallearningrate0 05 firstdecaysteps10 alpha0 1 tmul3 0 mmul4 0 namemycdr def npcosinedecayrestarts self step decaysteps tmul2 0 mmul1 0 alpha0 0 fac 1 0 while step decaysteps step decaysteps decaysteps tmul fac mmul completedfraction step decaysteps decay fac 0 5 1 0 math cosmath pi completedfraction return 1 0 alpha decay alpha def testdecayself numtrainingsteps 1000 initiallr 1 0 for step in range0 1500 250 decayedlr schedules cosinedecayrestarts initiallr numtrainingsteps expected self npcosinedecayrestartsstep numtrainingsteps self assertallclosedecayedlrstep expected 1e6 def testfloat64self numtrainingsteps 1000 initiallr np float641 0 for step in range0 1500 250 decayedlr schedules cosinedecayrestarts initiallr numtrainingsteps expected self npcosinedecayrestartsstep numtrainingsteps self assertallclosedecayedlrstep expected 1e6 def testalphaself numtrainingsteps 1000 initiallr 1 0 alpha 0 1 for step in range0 1500 250 decayedlr schedules cosinedecayrestarts initiallr numtrainingsteps alphaalpha expected self npcosinedecayrestarts step numtrainingsteps alphaalpha self assertallclosedecayedlrstep expected 1e6 def testmmulself numtrainingsteps 1000 initiallr 1 0 mmul 0 9 for step in range0 1500 250 decayedlr schedules cosinedecayrestarts initiallr numtrainingsteps mmulmmul expected self npcosinedecayrestarts step numtrainingsteps mmulmmul self assertallclosedecayedlrstep expected 1e6 def testtmulself numtrainingsteps 1000 initiallr 1 0 tmul 1 0 for step in range0 1500 250 decayedlr schedules cosinedecayrestarts initiallr numtrainingsteps tmultmul expected self npcosinedecayrestarts step numtrainingsteps tmultmul self assertallclosedecayedlrstep expected 1e6 no change to learning rate due to staircase decayed learning rate no change to learning rate decayed learning rate test casting boundaries from int32 to int64
import math

import numpy as np
import pytest

from keras import backend
from keras import layers
from keras import optimizers
from keras import testing
from keras.models import Sequential
from keras.optimizers import schedules


class TestFitLRSchedulesFlow(testing.TestCase):
    @pytest.mark.requires_trainable_backend
    def test_fit_lr_correctness(self):
        model = Sequential(
            [
                layers.Dense(
                    2, kernel_initializer="ones", bias_initializer="ones"
                )
            ]
        )
        optimizer = optimizers.Adam(
            learning_rate=schedules.ExponentialDecay(
                initial_learning_rate=0.05, decay_steps=1, decay_rate=0.9
            )
        )
        self.assertEqual(len(optimizer.variables), 1)
        self.assertEqual(optimizer.variables[0], 0)

        model.compile(optimizer=optimizer, loss="mse")
        x = np.arange(32).reshape((16, 2))
        y = np.arange(32).reshape((16, 2))
        history = model.fit(x, y, epochs=3, batch_size=4, shuffle=False)
        self.assertEqual(optimizer.variables[0], 4 * 3)
        self.assertAllClose(
            history.history["loss"],
            [230.79457092285156, 128.30319213867188, 79.33648681640625],
            rtol=5e-5,
        )


class ExponentialDecayTest(testing.TestCase):
    def test_config(self):
        self.run_class_serialization_test(
            schedules.ExponentialDecay(
                initial_learning_rate=0.05,
                decay_steps=10,
                decay_rate=0.96,
                staircase=True,
                name="my_ed",
            )
        )

    def test_continuous(self):
        step = 5
        decayed_lr = schedules.ExponentialDecay(0.05, 10, 0.96)
        expected = 0.05 * 0.96 ** (5.0 / 10.0)
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_staircase(self):
        step = backend.Variable(1.0)
        decayed_lr = schedules.ExponentialDecay(0.1, 3, 0.96, staircase=True)

        
        expected = 0.1
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

        expected = 0.1
        step.assign(2)
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

        
        expected = 0.1 * 0.96 ** (100 // 3)
        step.assign(100)
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_variables(self):
        step = backend.Variable(1.0)
        decayed_lr = schedules.ExponentialDecay(0.1, 3, 0.96, staircase=True)

        
        step.assign(1)
        self.assertAllClose(decayed_lr(step), 0.1, 1e-6)
        step.assign(2)
        self.assertAllClose(decayed_lr(step), 0.1, 1e-6)
        
        step.assign(100)
        expected = 0.1 * 0.96 ** (100 // 3)
        self.assertAllClose(decayed_lr(step), expected, 1e-6)


class PiecewiseConstantDecayTest(testing.TestCase):
    def test_config(self):
        self.run_class_serialization_test(
            schedules.PiecewiseConstantDecay(
                boundaries=[10, 20], values=[1, 2, 3], name="my_pcd"
            )
        )

    def test_piecewise_values(self):
        x = backend.Variable(-999.0)
        decayed_lr = schedules.PiecewiseConstantDecay(
            [100, 110, 120], [1.0, 0.1, 0.01, 0.001]
        )

        self.assertAllClose(decayed_lr(x), 1.0, 1e-6)
        x.assign(100)
        self.assertAllClose(decayed_lr(x), 1.0, 1e-6)
        x.assign(105)
        self.assertAllClose(decayed_lr(x), 0.1, 1e-6)
        x.assign(110)
        self.assertAllClose(decayed_lr(x), 0.1, 1e-6)
        x.assign(120)
        self.assertAllClose(decayed_lr(x), 0.01, 1e-6)
        x.assign(999)
        self.assertAllClose(decayed_lr(x), 0.001, 1e-6)

    def test_boundary_values(self):
        
        x_int64 = backend.Variable(0, dtype="int64", trainable=False)
        boundaries, values = [1, 2, 3], [0.4, 0.5, 0.6, 0.7]
        decayed_lr = schedules.PiecewiseConstantDecay(boundaries, values)

        self.assertAllClose(decayed_lr(x_int64), 0.4, 1e-6)
        x_int64.assign(1)
        self.assertAllClose(decayed_lr(x_int64), 0.4, 1e-6)
        x_int64.assign(2)
        self.assertAllClose(decayed_lr(x_int64), 0.5, 1e-6)
        x_int64.assign(3)
        self.assertAllClose(decayed_lr(x_int64), 0.6, 1e-6)
        x_int64.assign(4)
        self.assertAllClose(decayed_lr(x_int64), 0.7, 1e-6)


class LinearDecayTest(testing.TestCase):
    def test_config(self):
        self.run_class_serialization_test(
            schedules.PolynomialDecay(
                initial_learning_rate=0.1,
                decay_steps=100,
                end_learning_rate=0.005,
                power=1.0,
                cycle=False,
                name="my_ld",
            )
        )

    def test_halfway(self):
        step = 5
        lr = 0.05
        end_lr = 0.0
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr)
        expected = lr * 0.5
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_end(self):
        step = 10
        lr = 0.05
        end_lr = 0.001
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr)
        expected = end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_halfway_with_end(self):
        step = 5
        lr = 0.05
        end_lr = 0.001
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr)
        expected = (lr + end_lr) * 0.5
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_beyond_end(self):
        step = 15
        lr = 0.05
        end_lr = 0.001
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr)
        expected = end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_beyond_end_with_cycle(self):
        step = 15
        lr = 0.05
        end_lr = 0.001
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr, cycle=True)
        expected = (lr - end_lr) * 0.25 + end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)


class SqrtDecayTest(testing.TestCase):
    def test_halfway(self):
        step = 5
        lr = 0.05
        end_lr = 0.0
        power = 0.5
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr, power=power)
        expected = lr * 0.5**power
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_end(self):
        step = 10
        lr = 0.05
        end_lr = 0.001
        power = 0.5
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr, power=power)
        expected = end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_halfway_with_end(self):
        step = 5
        lr = 0.05
        end_lr = 0.001
        power = 0.5
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr, power=power)
        expected = (lr - end_lr) * 0.5**power + end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_beyond_end(self):
        step = 15
        lr = 0.05
        end_lr = 0.001
        power = 0.5
        decayed_lr = schedules.PolynomialDecay(lr, 10, end_lr, power=power)
        expected = end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_beyond_end_with_cycle(self):
        step = 15
        lr = 0.05
        end_lr = 0.001
        power = 0.5
        decayed_lr = schedules.PolynomialDecay(
            lr, 10, end_lr, power=power, cycle=True
        )
        expected = (lr - end_lr) * 0.25**power + end_lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_begin_with_cycle(self):
        lr = 0.001
        decay_steps = 10
        step = 0
        decayed_lr = schedules.PolynomialDecay(lr, decay_steps, cycle=True)
        expected = lr
        self.assertAllClose(decayed_lr(step), expected, 1e-6)


class InverseTimeDecayTest(testing.TestCase):
    def test_config(self):
        self.run_class_serialization_test(
            schedules.InverseTimeDecay(
                initial_learning_rate=0.05,
                decay_steps=10,
                decay_rate=0.96,
                staircase=True,
                name="my_itd",
            )
        )

    def test_decay(self):
        initial_lr = 0.1
        k = 10
        decay_rate = 0.96
        step = backend.Variable(0.0)
        decayed_lr = schedules.InverseTimeDecay(initial_lr, k, decay_rate)

        for i in range(k + 1):
            expected = initial_lr / (1 + i / k * decay_rate)
            self.assertAllClose(decayed_lr(step), expected, 1e-6)
            step.assign(step + 1)

    def test_staircase(self):
        initial_lr = 0.1
        k = 10
        decay_rate = 0.96
        step = backend.Variable(0.0)
        decayed_lr = schedules.InverseTimeDecay(
            initial_lr, k, decay_rate, staircase=True
        )

        for i in range(k + 1):
            expected = initial_lr / (1 + decay_rate * (i // k))
            self.assertAllClose(decayed_lr(step), expected, 1e-6)
            step.assign(step + 1)


class CosineDecayTest(testing.TestCase):
    def test_config(self):
        self.run_class_serialization_test(
            schedules.CosineDecay(
                initial_learning_rate=0.05,
                decay_steps=10,
                alpha=0.1,
                warmup_target=0.2,
                warmup_steps=2,
                name="my_cd",
            )
        )

    def np_cosine_decay(self, step, decay_steps, alpha=0.0):
        step = min(step, decay_steps)
        completed_fraction = step / decay_steps
        decay = 0.5 * (1.0 + math.cos(math.pi * completed_fraction))
        return (1.0 - alpha) * decay + alpha

    def test_decay(self):
        num_training_steps = 1000
        initial_lr = 1.0
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecay(initial_lr, num_training_steps)
            expected = self.np_cosine_decay(step, num_training_steps)
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def linear_warmup(self, step, warmup_steps, initial_lr, target_lr):
        completed_fraction = step / warmup_steps
        total_delta = target_lr - initial_lr
        return completed_fraction * total_delta

    def test_warmup(self):
        warmup_steps = 1500
        initial_lr = 0.0
        target_lr = 10.0
        for step in range(0, 1500, 250):
            lr = schedules.CosineDecay(
                initial_lr,
                10,
                warmup_target=target_lr,
                warmup_steps=warmup_steps,
            )
            expected = self.linear_warmup(
                step, warmup_steps, initial_lr, target_lr
            )
            self.assertAllClose(lr(step), expected)

    def test_alpha(self):
        num_training_steps = 1000
        initial_lr = 1.0
        alpha = 0.1
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecay(
                initial_lr, num_training_steps, alpha
            )
            expected = self.np_cosine_decay(step, num_training_steps, alpha)
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_float64(self):
        num_training_steps = 1000
        initial_lr = np.float64(1.0)
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecay(initial_lr, num_training_steps)
            expected = self.np_cosine_decay(step, num_training_steps)
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_warmup_decay(self):
        warmup_steps = 2000
        decay_steps = 1000
        initial_lr = 0.0
        target_lr = 10.0
        for step in range(0, 3000, 250):
            lr = schedules.CosineDecay(
                initial_lr,
                decay_steps,
                warmup_target=target_lr,
                warmup_steps=warmup_steps,
            )
            if step < warmup_steps + 1:
                expected = self.linear_warmup(
                    step, warmup_steps, initial_lr, target_lr
                )
            else:
                expected = target_lr * self.np_cosine_decay(
                    step - warmup_steps, decay_steps
                )
            self.assertAllClose(lr(step), expected)


class CosineDecayRestartsTest(testing.TestCase):
    def test_config(self):
        self.run_class_serialization_test(
            schedules.CosineDecayRestarts(
                initial_learning_rate=0.05,
                first_decay_steps=10,
                alpha=0.1,
                t_mul=3.0,
                m_mul=4.0,
                name="my_cdr",
            )
        )

    def np_cosine_decay_restarts(
        self, step, decay_steps, t_mul=2.0, m_mul=1.0, alpha=0.0
    ):
        fac = 1.0
        while step >= decay_steps:
            step -= decay_steps
            decay_steps *= t_mul
            fac *= m_mul

        completed_fraction = step / decay_steps
        decay = fac * 0.5 * (1.0 + math.cos(math.pi * completed_fraction))
        return (1.0 - alpha) * decay + alpha

    def test_decay(self):
        num_training_steps = 1000
        initial_lr = 1.0
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecayRestarts(
                initial_lr, num_training_steps
            )
            expected = self.np_cosine_decay_restarts(step, num_training_steps)
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_float64(self):
        num_training_steps = 1000
        initial_lr = np.float64(1.0)
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecayRestarts(
                initial_lr, num_training_steps
            )
            expected = self.np_cosine_decay_restarts(step, num_training_steps)
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_alpha(self):
        num_training_steps = 1000
        initial_lr = 1.0
        alpha = 0.1
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecayRestarts(
                initial_lr, num_training_steps, alpha=alpha
            )
            expected = self.np_cosine_decay_restarts(
                step, num_training_steps, alpha=alpha
            )
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_mmul(self):
        num_training_steps = 1000
        initial_lr = 1.0
        m_mul = 0.9
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecayRestarts(
                initial_lr, num_training_steps, m_mul=m_mul
            )
            expected = self.np_cosine_decay_restarts(
                step, num_training_steps, m_mul=m_mul
            )
            self.assertAllClose(decayed_lr(step), expected, 1e-6)

    def test_tmul(self):
        num_training_steps = 1000
        initial_lr = 1.0
        t_mul = 1.0
        for step in range(0, 1500, 250):
            decayed_lr = schedules.CosineDecayRestarts(
                initial_lr, num_training_steps, t_mul=t_mul
            )
            expected = self.np_cosine_decay_restarts(
                step, num_training_steps, t_mul=t_mul
            )
            self.assertAllClose(decayed_lr(step), expected, 1e-6)