various learning rate schedule functions import math from keras import ops from keras apiexport import kerasexport from keras saving import serializationlib kerasexportkeras optimizers schedules learningrateschedule class learningrateschedule def callself step raise notimplementederror flearning rate schedule self class name must override callself step def getconfigself raise notimplementederror flearning rate schedule self class name must override getconfig in order to be serializable classmethod def fromconfigcls config return clsconfig kerasexportkeras optimizers schedules exponentialdecay class exponentialdecaylearningrateschedule def init self initiallearningrate decaysteps decayrate staircasefalse nameexponentialdecay super init self initiallearningrate initiallearningrate self decaysteps decaysteps self decayrate decayrate self staircase staircase self name name if self decaysteps 0 raise valueerror argument decaysteps must be 0 freceived decaystepsself decaysteps def callself step with ops namescopeself name initiallearningrate ops converttotensor self initiallearningrate dtype initiallearningrate dtype decaysteps ops castself decaysteps dtype decayrate ops castself decayrate dtype globalsteprecomp ops caststep dtype p globalsteprecomp decaysteps if self staircase p ops floorp return ops multiplyinitiallearningrate ops powerdecayrate p def getconfigself return initiallearningrate self initiallearningrate decaysteps self decaysteps decayrate self decayrate staircase self staircase name self name kerasexportkeras optimizers schedules piecewiseconstantdecay class piecewiseconstantdecaylearningrateschedule later whenever we perform an optimization step we pass in the step def initself boundaries values namepiecewiseconstant super init if lenboundaries lenvalues 1 raise valueerror the length of boundaries should be 1 less than the length of fvalues received boundariesboundaries of length flenboundaries and valuesvalues fof length lenvalues self boundaries boundaries self values values self name name def callself step with ops namescopeself name boundaries ops converttotensorx for x in self boundaries values ops converttotensorx for x in self values step ops converttotensorstep for i b in enumerateboundaries if b dtype step dtype we cast the boundaries to have the same type as the step b ops castb step dtype boundariesi b resultdtype values0 dtype resultvalue ops array0 dtyperesultdtype for each range between boundaries we check whether the step is within that range cast the resulting boolean to a number and multiply the result by the corresponding value for the range taking the sum of these yields a piecewise constant function steplessthanfirstboundary ops cast step boundaries0 resultdtype resultvalue steplessthanfirstboundary values0 stepgreaterthanlastboundary ops cast step boundaries1 resultdtype resultvalue stepgreaterthanlastboundary values1 for low high value in zip boundaries 1 boundaries1 values1 1 need to bind v here can do this with lambda vv stepinrange ops cast step low step high resultdtype resultvalue stepinrange value return resultvalue def getconfigself return boundaries self boundaries values self values name self name kerasexportkeras optimizers schedules polynomialdecay class polynomialdecaylearningrateschedule def init self initiallearningrate decaysteps endlearningrate0 0001 power1 0 cyclefalse namepolynomialdecay super init self initiallearningrate initiallearningrate self decaysteps decaysteps self endlearningrate endlearningrate self power power self cycle cycle self name name if self decaysteps 0 raise valueerror argument decaysteps must be 0 freceived decaystepsself decaysteps def callself step with ops namescopeself name initiallearningrate ops converttotensor self initiallearningrate dtype initiallearningrate dtype endlearningrate ops castself endlearningrate dtype power ops castself power dtype globalsteprecomp ops caststep dtype decaystepsrecomp ops castself decaysteps dtype if self cycle find the first multiple of decaysteps that is bigger than globalstep if globalstep is zero set the multiplier to 1 multiplier ops where ops equalglobalsteprecomp 0 1 0 ops ceilglobalsteprecomp self decaysteps decaystepsrecomp ops multiply decaystepsrecomp multiplier else make sure that the globalstep used is not bigger than decaysteps globalsteprecomp ops minimum globalsteprecomp decaystepsrecomp p ops divideglobalsteprecomp decaystepsrecomp return ops add ops multiply initiallearningrate endlearningrate ops power1 p power endlearningrate def getconfigself return initiallearningrate self initiallearningrate decaysteps self decaysteps endlearningrate self endlearningrate power self power cycle self cycle name self name kerasexportkeras optimizers schedules inversetimedecay class inversetimedecaylearningrateschedule def init self initiallearningrate decaysteps decayrate staircasefalse nameinversetimedecay super init self initiallearningrate initiallearningrate self decaysteps decaysteps self decayrate decayrate self staircase staircase self name name if self decaysteps 0 raise valueerror argument decaysteps must be 0 freceived decaystepsself decaysteps def callself step with ops namescopeself name initiallearningrate ops converttotensor self initiallearningrate dtype initiallearningrate dtype decaysteps ops castself decaysteps dtype decayrate ops castself decayrate dtype globalsteprecomp ops caststep dtype p globalsteprecomp decaysteps if self staircase p ops floorp const ops castops array1 dtype denom ops addconst ops multiplydecayrate p return ops divideinitiallearningrate denom def getconfigself return initiallearningrate self initiallearningrate decaysteps self decaysteps decayrate self decayrate staircase self staircase name self name kerasexportkeras optimizers schedules cosinedecay class cosinedecaylearningrateschedule def init self initiallearningrate decaysteps alpha0 0 namecosinedecay warmuptargetnone warmupsteps0 super init self initiallearningrate initiallearningrate self decaysteps decaysteps self alpha alpha self name name self warmupsteps warmupsteps self warmuptarget warmuptarget if self decaysteps 0 raise valueerror argument decaysteps must be 0 freceived decaystepsself decaysteps def decayfunctionself step decaysteps decayfromlr dtype with ops namescopeself name completedfraction step decaysteps pi ops arraymath pi dtypedtype cosinedecayed 0 5 1 0 ops cospi completedfraction decayed 1 self alpha cosinedecayed self alpha return ops multiplydecayfromlr decayed def warmupfunction self step warmupsteps warmuptarget initiallearningrate with ops namescopeself name completedfraction step warmupsteps totalstepdelta warmuptarget initiallearningrate return totalstepdelta completedfraction initiallearningrate def callself step with ops namescopeself name initiallearningrate ops converttotensor self initiallearningrate dtype initiallearningrate dtype decaysteps ops castself decaysteps dtype globalsteprecomp ops caststep dtype if self warmuptarget is none globalsteprecomp ops minimum globalsteprecomp decaysteps return self decayfunction globalsteprecomp decaysteps initiallearningrate dtype warmuptarget ops castself warmuptarget dtype warmupsteps ops castself warmupsteps dtype globalsteprecomp ops minimum globalsteprecomp decaysteps warmupsteps return ops cond globalsteprecomp warmupsteps lambda self warmupfunction globalsteprecomp warmupsteps warmuptarget initiallearningrate lambda self decayfunction globalsteprecomp warmupsteps decaysteps warmuptarget dtype def getconfigself return initiallearningrate self initiallearningrate decaysteps self decaysteps alpha self alpha name self name warmuptarget self warmuptarget warmupsteps self warmupsteps kerasexportkeras optimizers schedules cosinedecayrestarts class cosinedecayrestartslearningrateschedule def init self initiallearningrate firstdecaysteps tmul2 0 mmul1 0 alpha0 0 namesgdrdecay super init self initiallearningrate initiallearningrate self firstdecaysteps firstdecaysteps self tmul tmul self mmul mmul self alpha alpha self name name if self firstdecaysteps 0 raise valueerror argument firstdecaysteps must be 0 freceived firstdecaystepsself firstdecaysteps def callself step with ops namescopeself name initiallearningrate ops converttotensor self initiallearningrate dtype initiallearningrate dtype firstdecaysteps ops castself firstdecaysteps dtype alpha ops castself alpha dtype tmul ops castself tmul dtype mmul ops castself mmul dtype globalsteprecomp ops caststep dtype completedfraction globalsteprecomp firstdecaysteps def computestepcompletedfraction geometricfalse ops log is sensitive to the precision of dtype so we need the additional casting serializes a learningrateschedule into a jsoncompatible dict args learningrateschedule the learningrateschedule object to serialize returns a jsonserializable dict representing the object s config example lrschedule keras optimizers schedules exponentialdecay 0 1 decaysteps100000 decayrate0 96 staircasetrue keras optimizers schedules serializelrschedule module keras optimizers schedules classname exponentialdecay config registeredname none instantiates a learningrateschedule object from a serialized form args config the serialized form of the learningrateschedule dictionary of the form classname str config dict customobjects a dictionary mapping class names or function names of custom nonkeras objects to classfunctions returns a learningrateschedule object example python configuration for polynomialdecay config classname polynomialdecay config cycle false decaysteps 10000 endlearningrate 0 01 initiallearningrate 0 1 name none power 0 5 lrschedule keras optimizers schedules deserializeconfig the learning rate schedule base class you can use a learning rate schedule to modulate how the learning rate of your optimizer changes over time several built in learning rate schedules are available such as keras optimizers schedules exponentialdecay or keras optimizers schedules piecewiseconstantdecay python lr_schedule keras optimizers schedules exponentialdecay initial_learning_rate 1e 2 decay_steps 10000 decay_rate 0 9 optimizer keras optimizers sgd learning_rate lr_schedule a learningrateschedule instance can be passed in as the learning_rate argument of any optimizer to implement your own schedule object you should implement the __call__ method which takes a step argument scalar integer tensor the current training step count like for any other keras object you can also optionally make your object serializable by implementing the get_config and from_config methods example python class mylrschedule keras optimizers schedules learningrateschedule def __init__ self initial_learning_rate self initial_learning_rate initial_learning_rate def __call__ self step return self initial_learning_rate step 1 optimizer keras optimizers sgd learning_rate mylrschedule 0 1 instantiates a learningrateschedule from its config args config output of get_config returns a learningrateschedule instance a learningrateschedule that uses an exponential decay schedule when training a model it is often useful to lower the learning rate as the training progresses this schedule applies an exponential decay function to an optimizer step given a provided initial learning rate the schedule is a 1 arg callable that produces a decayed learning rate when passed the current optimizer step this can be useful for changing the learning rate value across different invocations of optimizer functions it is computed as python def decayed_learning_rate step return initial_learning_rate decay_rate step decay_steps if the argument staircase is true then step decay_steps is an integer division and the decayed learning rate follows a staircase function you can pass this schedule directly into a keras optimizers optimizer as the learning rate example when fitting a keras model decay every 100000 steps with a base of 0 96 python initial_learning_rate 0 1 lr_schedule keras optimizers schedules exponentialdecay initial_learning_rate decay_steps 100000 decay_rate 0 96 staircase true model compile optimizer keras optimizers sgd learning_rate lr_schedule loss sparse_categorical_crossentropy metrics accuracy model fit data labels epochs 5 the learning rate schedule is also serializable and deserializable using keras optimizers schedules serialize and keras optimizers schedules deserialize args initial_learning_rate a python float the initial learning rate decay_steps a python integer must be positive see the decay computation above decay_rate a python float the decay rate staircase boolean if true decay the learning rate at discrete intervals name string optional name of the operation defaults to exponentialdecay returns a 1 arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate a scalar tensor of the same type as initial_learning_rate a learningrateschedule that uses a piecewise constant decay schedule the function returns a 1 arg callable to compute the piecewise constant when passed the current optimizer step this can be useful for changing the learning rate value across different invocations of optimizer functions example use a learning rate that s 1 0 for the first 100001 steps 0 5 for the next 10000 steps and 0 1 for any additional steps python step ops array 0 boundaries 100000 110000 values 1 0 0 5 0 1 learning_rate_fn keras optimizers schedules piecewiseconstantdecay boundaries values later whenever we perform an optimization step we pass in the step learning_rate learning_rate_fn step you can pass this schedule directly into a keras optimizers optimizer as the learning rate the learning rate schedule is also serializable and deserializable using keras optimizers schedules serialize and keras optimizers schedules deserialize args boundaries a list of python numbers with strictly increasing entries and with all elements having the same type as the optimizer step values a list of python numbers that specifies the values for the intervals defined by boundaries it should have one more element than boundaries and all elements should have the same type name a string optional name of the operation defaults to piecewiseconstant returns a 1 arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate a scalar tensor of the same type as the boundary tensors the output of the 1 arg function that takes the step is values 0 when step boundaries 0 values 1 when step boundaries 0 and step boundaries 1 and values 1 when step boundaries 1 raises valueerror if the number of elements in the boundaries and values lists do not match we cast the boundaries to have the same type as the step for each range between boundaries we check whether the step is within that range cast the resulting boolean to a number and multiply the result by the corresponding value for the range taking the sum of these yields a piecewise constant function need to bind v here can do this with lambda v v a learningrateschedule that uses a polynomial decay schedule it is commonly observed that a monotonically decreasing learning rate whose degree of change is carefully chosen results in a better performing model this schedule applies a polynomial decay function to an optimizer step given a provided initial_learning_rate to reach an end_learning_rate in the given decay_steps it requires a step value to compute the decayed learning rate you can just pass a backend variable that you increment at each training step the schedule is a 1 arg callable that produces a decayed learning rate when passed the current optimizer step this can be useful for changing the learning rate value across different invocations of optimizer functions it is computed as python def decayed_learning_rate step step min step decay_steps return initial_learning_rate end_learning_rate 1 step decay_steps power end_learning_rate if cycle is true then a multiple of decay_steps is used the first one that is bigger than step python def decayed_learning_rate step decay_steps decay_steps ceil step decay_steps return initial_learning_rate end_learning_rate 1 step decay_steps power end_learning_rate you can pass this schedule directly into a keras optimizers optimizer as the learning rate example fit a model while decaying from 0 1 to 0 01 in 10000 steps using sqrt i e power 0 5 python starter_learning_rate 0 1 end_learning_rate 0 01 decay_steps 10000 learning_rate_fn keras optimizers schedules polynomialdecay starter_learning_rate decay_steps end_learning_rate power 0 5 model compile optimizer keras optimizers sgd learning_rate learning_rate_fn loss sparse_categorical_crossentropy metrics accuracy model fit data labels epochs 5 the learning rate schedule is also serializable and deserializable using keras optimizers schedules serialize and keras optimizers schedules deserialize args initial_learning_rate a python float the initial learning rate decay_steps a python integer must be positive see the decay computation above end_learning_rate a python float the minimal end learning rate power a python float the power of the polynomial defaults to 1 0 cycle a boolean whether it should cycle beyond decay_steps name string optional name of the operation defaults to polynomialdecay returns a 1 arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate a scalar tensor of the same type as initial_learning_rate find the first multiple of decay_steps that is bigger than global_step if global_step is zero set the multiplier to 1 make sure that the global_step used is not bigger than decay_steps a learningrateschedule that uses an inverse time decay schedule when training a model it is often useful to lower the learning rate as the training progresses this schedule applies the inverse decay function to an optimizer step given a provided initial learning rate it requires a step value to compute the decayed learning rate you can just pass a backend variable that you increment at each training step the schedule is a 1 arg callable that produces a decayed learning rate when passed the current optimizer step this can be useful for changing the learning rate value across different invocations of optimizer functions it is computed as python def decayed_learning_rate step return initial_learning_rate 1 decay_rate step decay_step or if staircase is true as python def decayed_learning_rate step return initial_learning_rate 1 decay_rate floor step decay_step you can pass this schedule directly into a keras optimizers optimizer as the learning rate example fit a keras model when decaying 1 t with a rate of 0 5 python initial_learning_rate 0 1 decay_steps 1 0 decay_rate 0 5 learning_rate_fn keras optimizers schedules inversetimedecay initial_learning_rate decay_steps decay_rate model compile optimizer keras optimizers sgd learning_rate learning_rate_fn loss sparse_categorical_crossentropy metrics accuracy model fit data labels epochs 5 args initial_learning_rate a python float the initial learning rate decay_steps how often to apply decay decay_rate a python number the decay rate staircase whether to apply decay in a discrete staircase as o pposed to continuous fashion name string optional name of the operation defaults to inversetimedecay returns a 1 arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate a scalar tensor of the same type as initial_learning_rate a learningrateschedule that uses a cosine decay with optional warmup see loshchilov hutter iclr2016 https arxiv org abs 1608 03983 sgdr stochastic gradient descent with warm restarts for the idea of a linear warmup of our learning rate see goyal et al https arxiv org pdf 1706 02677 pdf when we begin training a model we often want an initial increase in our learning rate followed by a decay if warmup_target is an int this schedule applies a linear increase per optimizer step to our learning rate from initial_learning_rate to warmup_target for a duration of warmup_steps afterwards it applies a cosine decay function taking our learning rate from warmup_target to alpha for a duration of decay_steps if warmup_target is none we skip warmup and our decay will take our learning rate from initial_learning_rate to alpha it requires a step value to compute the learning rate you can just pass a backend variable that you increment at each training step the schedule is a 1 arg callable that produces a warmup followed by a decayed learning rate when passed the current optimizer step this can be useful for changing the learning rate value across different invocations of optimizer functions our warmup is computed as python def warmup_learning_rate step completed_fraction step warmup_steps total_delta target_warmup initial_learning_rate return completed_fraction total_delta and our decay is computed as python if warmup_target is none initial_decay_lr initial_learning_rate else initial_decay_lr warmup_target def decayed_learning_rate step step min step decay_steps cosine_decay 0 5 1 cos pi step decay_steps decayed 1 alpha cosine_decay alpha return initial_decay_lr decayed example usage without warmup python decay_steps 1000 initial_learning_rate 0 1 lr_decayed_fn keras optimizers schedules cosinedecay initial_learning_rate decay_steps example usage with warmup python decay_steps 1000 initial_learning_rate 0 warmup_steps 1000 target_learning_rate 0 1 lr_warmup_decayed_fn keras optimizers schedules cosinedecay initial_learning_rate decay_steps warmup_target target_learning_rate warmup_steps warmup_steps you can pass this schedule directly into a keras optimizers optimizer as the learning rate the learning rate schedule is also serializable and deserializable using keras optimizers schedules serialize and keras optimizers schedules deserialize args initial_learning_rate a python float the initial learning rate decay_steps a python int number of steps to decay over alpha a python float minimum learning rate value for decay as a fraction of initial_learning_rate name string optional name of the operation defaults to cosinedecay warmup_target a python float the target learning rate for our warmup phase will cast to the initial_learning_rate datatype setting to none will skip warmup and begins decay phase from initial_learning_rate otherwise scheduler will warmup from initial_learning_rate to warmup_target warmup_steps a python int number of steps to warmup over returns a 1 arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate a scalar tensor of the same type as initial_learning_rate a learningrateschedule that uses a cosine decay schedule with restarts see loshchilov hutter iclr2016 https arxiv org abs 1608 03983 sgdr stochastic gradient descent with warm restarts when training a model it is often useful to lower the learning rate as the training progresses this schedule applies a cosine decay function with restarts to an optimizer step given a provided initial learning rate it requires a step value to compute the decayed learning rate you can just pass a backend variable that you increment at each training step the schedule is a 1 arg callable that produces a decayed learning rate when passed the current optimizer step this can be useful for changing the learning rate value across different invocations of optimizer functions the learning rate multiplier first decays from 1 to alpha for first_decay_steps steps then a warm restart is performed each new warm restart runs for t_mul times more steps and with m_mul times initial learning rate as the new learning rate example usage python first_decay_steps 1000 lr_decayed_fn keras optimizers schedules cosinedecayrestarts initial_learning_rate first_decay_steps you can pass this schedule directly into a keras optimizers optimizer as the learning rate the learning rate schedule is also serializable and deserializable using keras optimizers schedules serialize and keras optimizers schedules deserialize args initial_learning_rate a python float the initial learning rate first_decay_steps a python integer number of steps to decay over t_mul a python float used to derive the number of iterations in the i th period m_mul a python float used to derive the initial learning rate of the i th period alpha a python float minimum learning rate value as a fraction of the initial_learning_rate name string optional name of the operation defaults to sgdrdecay returns a 1 arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate a scalar tensor of the same type as initial_learning_rate helper for cond operation ops log is sensitive to the precision of dtype so we need the additional casting serializes a learningrateschedule into a json compatible dict args learning_rate_schedule the learningrateschedule object to serialize returns a json serializable dict representing the object s config example lr_schedule keras optimizers schedules exponentialdecay 0 1 decay_steps 100000 decay_rate 0 96 staircase true keras optimizers schedules serialize lr_schedule module keras optimizers schedules class_name exponentialdecay config registered_name none instantiates a learningrateschedule object from a serialized form args config the serialized form of the learningrateschedule dictionary of the form class_name str config dict custom_objects a dictionary mapping class names or function names of custom non keras objects to class functions returns a learningrateschedule object example python configuration for polynomialdecay config class_name polynomialdecay config cycle false decay_steps 10000 end_learning_rate 0 01 initial_learning_rate 0 1 name none power 0 5 lr_schedule keras optimizers schedules deserialize config
import math

from keras import ops
from keras.api_export import keras_export
from keras.saving import serialization_lib


@keras_export("keras.optimizers.schedules.LearningRateSchedule")
class LearningRateSchedule:
    

    def __call__(self, step):
        raise NotImplementedError(
            f"Learning rate schedule '{self.__class__.__name__}' "
            "must override `__call__(self, step)`."
        )

    def get_config(self):
        raise NotImplementedError(
            f"Learning rate schedule '{self.__class__.__name__}' "
            "must override `get_config()` in order to be serializable."
        )

    @classmethod
    def from_config(cls, config):
        
        return cls(**config)


@keras_export("keras.optimizers.schedules.ExponentialDecay")
class ExponentialDecay(LearningRateSchedule):
    

    def __init__(
        self,
        initial_learning_rate,
        decay_steps,
        decay_rate,
        staircase=False,
        name="ExponentialDecay",
    ):
        super().__init__()
        self.initial_learning_rate = initial_learning_rate
        self.decay_steps = decay_steps
        self.decay_rate = decay_rate
        self.staircase = staircase
        self.name = name

        if self.decay_steps <= 0:
            raise ValueError(
                "Argument `decay_steps` must be > 0. "
                f"Received: decay_steps={self.decay_steps}"
            )

    def __call__(self, step):
        with ops.name_scope(self.name):
            initial_learning_rate = ops.convert_to_tensor(
                self.initial_learning_rate
            )
            dtype = initial_learning_rate.dtype
            decay_steps = ops.cast(self.decay_steps, dtype)
            decay_rate = ops.cast(self.decay_rate, dtype)

            global_step_recomp = ops.cast(step, dtype)
            p = global_step_recomp / decay_steps
            if self.staircase:
                p = ops.floor(p)
            return ops.multiply(initial_learning_rate, ops.power(decay_rate, p))

    def get_config(self):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_steps": self.decay_steps,
            "decay_rate": self.decay_rate,
            "staircase": self.staircase,
            "name": self.name,
        }


@keras_export("keras.optimizers.schedules.PiecewiseConstantDecay")
class PiecewiseConstantDecay(LearningRateSchedule):
    

    def __init__(self, boundaries, values, name="PiecewiseConstant"):
        super().__init__()

        if len(boundaries) != len(values) - 1:
            raise ValueError(
                "The length of boundaries should be 1 less than the length of "
                f"values. Received: boundaries={boundaries} of length "
                f"{len(boundaries)}, and values={values} "
                f"of length {len(values)}."
            )

        self.boundaries = boundaries
        self.values = values
        self.name = name

    def __call__(self, step):
        with ops.name_scope(self.name):
            boundaries = [ops.convert_to_tensor(x) for x in self.boundaries]
            values = [ops.convert_to_tensor(x) for x in self.values]
            step = ops.convert_to_tensor(step)

            for i, b in enumerate(boundaries):
                if b.dtype != step.dtype:
                    
                    b = ops.cast(b, step.dtype)
                    boundaries[i] = b

            result_dtype = values[0].dtype
            result_value = ops.array(0, dtype=result_dtype)

            
            
            
            
            step_less_than_first_boundary = ops.cast(
                step <= boundaries[0], result_dtype
            )
            result_value += step_less_than_first_boundary * values[0]

            step_greater_than_last_boundary = ops.cast(
                step > boundaries[-1], result_dtype
            )
            result_value += step_greater_than_last_boundary * values[-1]

            for low, high, value in zip(
                boundaries[:-1], boundaries[1:], values[1:-1]
            ):
                
                step_in_range = ops.cast(
                    (step > low) & (step <= high), result_dtype
                )
                result_value += step_in_range * value

            return result_value

    def get_config(self):
        return {
            "boundaries": self.boundaries,
            "values": self.values,
            "name": self.name,
        }


@keras_export("keras.optimizers.schedules.PolynomialDecay")
class PolynomialDecay(LearningRateSchedule):
    

    def __init__(
        self,
        initial_learning_rate,
        decay_steps,
        end_learning_rate=0.0001,
        power=1.0,
        cycle=False,
        name="PolynomialDecay",
    ):
        super().__init__()

        self.initial_learning_rate = initial_learning_rate
        self.decay_steps = decay_steps
        self.end_learning_rate = end_learning_rate
        self.power = power
        self.cycle = cycle
        self.name = name

        if self.decay_steps <= 0:
            raise ValueError(
                "Argument `decay_steps` must be > 0. "
                f"Received: decay_steps={self.decay_steps}"
            )

    def __call__(self, step):
        with ops.name_scope(self.name):
            initial_learning_rate = ops.convert_to_tensor(
                self.initial_learning_rate
            )
            dtype = initial_learning_rate.dtype
            end_learning_rate = ops.cast(self.end_learning_rate, dtype)
            power = ops.cast(self.power, dtype)

            global_step_recomp = ops.cast(step, dtype)
            decay_steps_recomp = ops.cast(self.decay_steps, dtype)
            if self.cycle:
                
                
                multiplier = ops.where(
                    ops.equal(global_step_recomp, 0),
                    1.0,
                    ops.ceil(global_step_recomp / self.decay_steps),
                )
                decay_steps_recomp = ops.multiply(
                    decay_steps_recomp, multiplier
                )
            else:
                
                
                global_step_recomp = ops.minimum(
                    global_step_recomp, decay_steps_recomp
                )

            p = ops.divide(global_step_recomp, decay_steps_recomp)
            return ops.add(
                ops.multiply(
                    initial_learning_rate - end_learning_rate,
                    ops.power(1 - p, power),
                ),
                end_learning_rate,
            )

    def get_config(self):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_steps": self.decay_steps,
            "end_learning_rate": self.end_learning_rate,
            "power": self.power,
            "cycle": self.cycle,
            "name": self.name,
        }


@keras_export("keras.optimizers.schedules.InverseTimeDecay")
class InverseTimeDecay(LearningRateSchedule):
    

    def __init__(
        self,
        initial_learning_rate,
        decay_steps,
        decay_rate,
        staircase=False,
        name="InverseTimeDecay",
    ):
        super().__init__()

        self.initial_learning_rate = initial_learning_rate
        self.decay_steps = decay_steps
        self.decay_rate = decay_rate
        self.staircase = staircase
        self.name = name

        if self.decay_steps <= 0:
            raise ValueError(
                "Argument `decay_steps` must be > 0. "
                f"Received: decay_steps={self.decay_steps}"
            )

    def __call__(self, step):
        with ops.name_scope(self.name):
            initial_learning_rate = ops.convert_to_tensor(
                self.initial_learning_rate
            )
            dtype = initial_learning_rate.dtype
            decay_steps = ops.cast(self.decay_steps, dtype)
            decay_rate = ops.cast(self.decay_rate, dtype)

            global_step_recomp = ops.cast(step, dtype)
            p = global_step_recomp / decay_steps
            if self.staircase:
                p = ops.floor(p)
            const = ops.cast(ops.array(1), dtype)
            denom = ops.add(const, ops.multiply(decay_rate, p))
            return ops.divide(initial_learning_rate, denom)

    def get_config(self):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_steps": self.decay_steps,
            "decay_rate": self.decay_rate,
            "staircase": self.staircase,
            "name": self.name,
        }


@keras_export("keras.optimizers.schedules.CosineDecay")
class CosineDecay(LearningRateSchedule):
    

    def __init__(
        self,
        initial_learning_rate,
        decay_steps,
        alpha=0.0,
        name="CosineDecay",
        warmup_target=None,
        warmup_steps=0,
    ):
        super().__init__()

        self.initial_learning_rate = initial_learning_rate
        self.decay_steps = decay_steps
        self.alpha = alpha
        self.name = name
        self.warmup_steps = warmup_steps
        self.warmup_target = warmup_target

        if self.decay_steps <= 0:
            raise ValueError(
                "Argument `decay_steps` must be > 0. "
                f"Received: decay_steps={self.decay_steps}"
            )

    def _decay_function(self, step, decay_steps, decay_from_lr, dtype):
        with ops.name_scope(self.name):
            completed_fraction = step / decay_steps
            pi = ops.array(math.pi, dtype=dtype)
            cosine_decayed = 0.5 * (1.0 + ops.cos(pi * completed_fraction))
            decayed = (1 - self.alpha) * cosine_decayed + self.alpha
            return ops.multiply(decay_from_lr, decayed)

    def _warmup_function(
        self, step, warmup_steps, warmup_target, initial_learning_rate
    ):
        with ops.name_scope(self.name):
            completed_fraction = step / warmup_steps
            total_step_delta = warmup_target - initial_learning_rate
            return total_step_delta * completed_fraction + initial_learning_rate

    def __call__(self, step):
        with ops.name_scope(self.name):
            initial_learning_rate = ops.convert_to_tensor(
                self.initial_learning_rate
            )
            dtype = initial_learning_rate.dtype
            decay_steps = ops.cast(self.decay_steps, dtype)
            global_step_recomp = ops.cast(step, dtype)

            if self.warmup_target is None:
                global_step_recomp = ops.minimum(
                    global_step_recomp, decay_steps
                )
                return self._decay_function(
                    global_step_recomp,
                    decay_steps,
                    initial_learning_rate,
                    dtype,
                )

            warmup_target = ops.cast(self.warmup_target, dtype)
            warmup_steps = ops.cast(self.warmup_steps, dtype)

            global_step_recomp = ops.minimum(
                global_step_recomp, decay_steps + warmup_steps
            )

            return ops.cond(
                global_step_recomp < warmup_steps,
                lambda: self._warmup_function(
                    global_step_recomp,
                    warmup_steps,
                    warmup_target,
                    initial_learning_rate,
                ),
                lambda: self._decay_function(
                    global_step_recomp - warmup_steps,
                    decay_steps,
                    warmup_target,
                    dtype,
                ),
            )

    def get_config(self):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_steps": self.decay_steps,
            "alpha": self.alpha,
            "name": self.name,
            "warmup_target": self.warmup_target,
            "warmup_steps": self.warmup_steps,
        }


@keras_export("keras.optimizers.schedules.CosineDecayRestarts")
class CosineDecayRestarts(LearningRateSchedule):
    

    def __init__(
        self,
        initial_learning_rate,
        first_decay_steps,
        t_mul=2.0,
        m_mul=1.0,
        alpha=0.0,
        name="SGDRDecay",
    ):
        super().__init__()

        self.initial_learning_rate = initial_learning_rate
        self.first_decay_steps = first_decay_steps
        self._t_mul = t_mul
        self._m_mul = m_mul
        self.alpha = alpha
        self.name = name

        if self.first_decay_steps <= 0:
            raise ValueError(
                "Argument `first_decay_steps` must be > 0. "
                f"Received: first_decay_steps={self.first_decay_steps}"
            )

    def __call__(self, step):
        with ops.name_scope(self.name):
            initial_learning_rate = ops.convert_to_tensor(
                self.initial_learning_rate
            )
            dtype = initial_learning_rate.dtype
            first_decay_steps = ops.cast(self.first_decay_steps, dtype)
            alpha = ops.cast(self.alpha, dtype)
            t_mul = ops.cast(self._t_mul, dtype)
            m_mul = ops.cast(self._m_mul, dtype)

            global_step_recomp = ops.cast(step, dtype)
            completed_fraction = global_step_recomp / first_decay_steps

            def compute_step(completed_fraction, geometric=False):
                
                if geometric:
                    
                    
                    i_restart = ops.floor(
                        ops.log(
                            ops.cast(
                                1.0 - completed_fraction * (1.0 - t_mul), dtype
                            )
                        )
                        / ops.log(t_mul)
                    )

                    sum_r = (1.0 - t_mul**i_restart) / (1.0 - t_mul)
                    completed_fraction = (
                        completed_fraction - sum_r
                    ) / t_mul**i_restart

                else:
                    i_restart = ops.floor(completed_fraction)
                    completed_fraction -= i_restart

                return i_restart, completed_fraction

            i_restart, completed_fraction = ops.cond(
                ops.equal(t_mul, 1.0),
                lambda: compute_step(completed_fraction, geometric=False),
                lambda: compute_step(completed_fraction, geometric=True),
            )

            m_fac = m_mul**i_restart
            cosine_decayed = (
                0.5
                * m_fac
                * (
                    1.0
                    + ops.cos(
                        ops.array(math.pi, dtype=dtype) * completed_fraction
                    )
                )
            )
            decayed = (1 - alpha) * cosine_decayed + alpha

            return ops.multiply(initial_learning_rate, decayed)

    def get_config(self):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "first_decay_steps": self.first_decay_steps,
            "t_mul": self._t_mul,
            "m_mul": self._m_mul,
            "alpha": self.alpha,
            "name": self.name,
        }


@keras_export("keras.optimizers.schedules.serialize")
def serialize(learning_rate_schedule):
    
    return serialization_lib.serialize_keras_object(learning_rate_schedule)


@keras_export("keras.optimizers.schedules.deserialize")
def deserialize(config, custom_objects=None):
    
    return serialization_lib.deserialize_keras_object(
        config,
        module_objects=globals(),
        custom_objects=custom_objects,
        printable_module_name="decay",
    )