optimizer that implements the adamax algorithm adamax a variant of adam based on the infinity norm is a firstorder gradientbased optimization method due to its capability of adjusting the learning rate based on data characteristics it is suited to learn timevariant process e g speech data with dynamically changed noise conditions default parameters follow those provided in the paper see references below initialization python m 0 initialize initial 1st moment vector u 0 initialize the exponentially weighted infinity norm t 0 initialize timestep the update rule for parameter w with gradient g is described at the end of section 7 1 of the paper see the referenece section python t 1 m beta1 m 1 beta g u maxbeta2 u absg currentlr learningrate 1 beta1 t w w currentlr m u epsilon args learningrate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta1 a float value or a constant float tensor the exponential decay rate for the 1st moment estimates beta2 a float value or a constant float tensor the exponential decay rate for the exponentially weighted infinity norm epsilon a small constant for numerical stability baseoptimizerkeywordargs reference kingma et al 2014http arxiv orgabs1412 6980 initialize optimizer variables adamax optimizer has 2 types of variables momentums denoted as m exponentially weighted infinity norm denoted as u args varlist list of model variables to build adamax variables on update step given gradient and the associated model variable lr ops castlearningrate variable dtype gradient ops castgradient variable dtype localstep ops castself iterations 1 variable dtype beta1power ops power ops castself beta1 variable dtype localstep m self mself getvariableindexvariable u self uself getvariableindexvariable self assignadd m ops multiplyops subtractgradient m 1 self beta1 self assign u ops maximumops multiplyself beta2 u ops absgradient self assignsub variable ops divide ops multiplylr m ops multiply1 beta1power ops addu self epsilon def getconfigself config super getconfig config update beta1 self beta1 beta2 self beta2 epsilon self epsilon return config adamax doc adamax doc replace baseoptimizerkeywordargs optimizer baseoptimizerkeywordargs optimizer that implements the adamax algorithm adamax a variant of adam based on the infinity norm is a first order gradient based optimization method due to its capability of adjusting the learning rate based on data characteristics it is suited to learn time variant process e g speech data with dynamically changed noise conditions default parameters follow those provided in the paper see references below initialization python m 0 initialize initial 1st moment vector u 0 initialize the exponentially weighted infinity norm t 0 initialize timestep the update rule for parameter w with gradient g is described at the end of section 7 1 of the paper see the referenece section python t 1 m beta1 m 1 beta g u max beta2 u abs g current_lr learning_rate 1 beta1 t w w current_lr m u epsilon args learning_rate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta_1 a float value or a constant float tensor the exponential decay rate for the 1st moment estimates beta_2 a float value or a constant float tensor the exponential decay rate for the exponentially weighted infinity norm epsilon a small constant for numerical stability base_optimizer_keyword_args reference kingma et al 2014 http arxiv org abs 1412 6980 initialize optimizer variables adamax optimizer has 2 types of variables momentums denoted as m exponentially weighted infinity norm denoted as u args var_list list of model variables to build adamax variables on update step given gradient and the associated model variable
from keras import ops
from keras.api_export import keras_export
from keras.optimizers import optimizer


@keras_export(["keras.optimizers.Adamax"])
class Adamax(optimizer.Optimizer):
    

    def __init__(
        self,
        learning_rate=0.001,
        beta_1=0.9,
        beta_2=0.999,
        epsilon=1e-7,
        weight_decay=None,
        clipnorm=None,
        clipvalue=None,
        global_clipnorm=None,
        use_ema=False,
        ema_momentum=0.99,
        ema_overwrite_frequency=None,
        name="adamax",
        **kwargs,
    ):
        super().__init__(
            learning_rate=learning_rate,
            name=name,
            weight_decay=weight_decay,
            clipnorm=clipnorm,
            clipvalue=clipvalue,
            global_clipnorm=global_clipnorm,
            use_ema=use_ema,
            ema_momentum=ema_momentum,
            ema_overwrite_frequency=ema_overwrite_frequency,
            **kwargs,
        )
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.epsilon = epsilon

    def build(self, var_list):
        
        if self.built:
            return
        super().build(var_list)
        self._m = []
        self._u = []
        for var in var_list:
            self._m.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="momentum"
                )
            )
            self._u.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="norm"
                )
            )

    def update_step(self, gradient, variable, learning_rate):
        
        lr = ops.cast(learning_rate, variable.dtype)
        gradient = ops.cast(gradient, variable.dtype)
        local_step = ops.cast(self.iterations + 1, variable.dtype)
        beta_1_power = ops.power(
            ops.cast(self.beta_1, variable.dtype), local_step
        )

        m = self._m[self._get_variable_index(variable)]
        u = self._u[self._get_variable_index(variable)]

        self.assign_add(
            m, ops.multiply(ops.subtract(gradient, m), (1 - self.beta_1))
        )
        self.assign(
            u, ops.maximum(ops.multiply(self.beta_2, u), ops.abs(gradient))
        )
        self.assign_sub(
            variable,
            ops.divide(
                ops.multiply(lr, m),
                ops.multiply((1 - beta_1_power), ops.add(u, self.epsilon)),
            ),
        )

    def get_config(self):
        config = super().get_config()

        config.update(
            {
                "beta_1": self.beta_1,
                "beta_2": self.beta_2,
                "epsilon": self.epsilon,
            }
        )
        return config


Adamax.__doc__ = Adamax.__doc__.replace(
    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
)