optimizer that implements the nadam algorithm much like adam is essentially rmsprop with momentum nadam is adam with nesterov momentum 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 or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 1st moment estimates defaults to 0 9 beta2 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 2nd moment estimates defaults to 0 999 epsilon a small constant for numerical stability this epsilon is epsilon hat in the kingma and ba paper in the formula just before section 2 1 not the epsilon in algorithm 1 of the paper defaults to 1e7 baseoptimizerkeywordargs reference dozat 2015http cs229 stanford eduproj2015054report pdf initialize optimizer variables nadam optimizer has 2 types of variables momentums and velocities args varlist list of model variables to build nadam variables on update step given gradient and the associated model variable vardtype variable dtype lr ops castlearningrate vardtype gradient ops castgradient vardtype localstep ops castself iterations 1 vardtype nextstep ops castself iterations 2 vardtype decay ops cast0 96 vardtype beta1 ops castself beta1 vardtype beta2 ops castself beta2 vardtype ut beta1 1 0 0 5 ops powerdecay localstep ut1 beta1 1 0 0 5 ops powerdecay nextstep uproductt ops castself uproduct vardtype uproductt1 uproductt ut1 beta2power ops powerbeta2 localstep m self momentumsself getvariableindexvariable v self velocitiesself getvariableindexvariable self assignadd m ops multiplyops subtractgradient m 1 beta1 self assignadd v ops multiplyops subtractops squaregradient v 1 beta2 mhat ops add ops divideops multiplyut1 m 1 uproductt1 ops divideops multiply1 ut gradient 1 uproductt vhat ops dividev 1 beta2power self assignsub variable ops divide ops multiplymhat lr ops addops sqrtvhat self epsilon def getconfigself config super getconfig config update beta1 self beta1 beta2 self beta2 epsilon self epsilon return config nadam doc nadam doc replace baseoptimizerkeywordargs optimizer baseoptimizerkeywordargs optimizer that implements the nadam algorithm much like adam is essentially rmsprop with momentum nadam is adam with nesterov momentum 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 or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 1st moment estimates defaults to 0 9 beta_2 a float value or a constant float tensor or a callable that takes no arguments and returns the actual value to use the exponential decay rate for the 2nd moment estimates defaults to 0 999 epsilon a small constant for numerical stability this epsilon is epsilon hat in the kingma and ba paper in the formula just before section 2 1 not the epsilon in algorithm 1 of the paper defaults to 1e 7 base_optimizer_keyword_args reference dozat 2015 http cs229 stanford edu proj2015 054_report pdf initialize optimizer variables nadam optimizer has 2 types of variables momentums and velocities args var_list list of model variables to build nadam variables on update step given gradient and the associated model variable
from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.optimizers import optimizer


@keras_export(["keras.optimizers.Nadam"])
class Nadam(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="nadam",
        **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
        if var_list:
            dtype = var_list[0].dtype
        else:
            dtype = backend.floatx()
        super().build(var_list)
        self._momentums = []
        self._velocities = []
        self._u_product = backend.Variable(1.0, dtype=dtype)

        for var in var_list:
            self._momentums.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="momentum"
                )
            )
            self._velocities.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="velocity"
                )
            )

    def _internal_apply_gradients(self, grads_and_vars):
        dtype = self._u_product.dtype
        self.assign(
            self._u_product,
            self._u_product
            * self.beta_1
            * (
                1.0
                - 0.5 * ops.power(0.96, ops.cast(self.iterations + 1, dtype))
            ),
        )
        super()._internal_apply_gradients(grads_and_vars)

    def update_step(self, gradient, variable, learning_rate):
        
        var_dtype = variable.dtype
        lr = ops.cast(learning_rate, var_dtype)
        gradient = ops.cast(gradient, var_dtype)

        local_step = ops.cast(self.iterations + 1, var_dtype)
        next_step = ops.cast(self.iterations + 2, var_dtype)
        decay = ops.cast(0.96, var_dtype)
        beta_1 = ops.cast(self.beta_1, var_dtype)
        beta_2 = ops.cast(self.beta_2, var_dtype)
        u_t = beta_1 * (1.0 - 0.5 * (ops.power(decay, local_step)))
        u_t_1 = beta_1 * (1.0 - 0.5 * (ops.power(decay, next_step)))
        u_product_t = ops.cast(self._u_product, var_dtype)

        u_product_t_1 = u_product_t * u_t_1
        beta_2_power = ops.power(beta_2, local_step)

        m = self._momentums[self._get_variable_index(variable)]
        v = self._velocities[self._get_variable_index(variable)]

        self.assign_add(
            m, ops.multiply(ops.subtract(gradient, m), (1 - beta_1))
        )
        self.assign_add(
            v, ops.multiply(ops.subtract(ops.square(gradient), v), (1 - beta_2))
        )
        m_hat = ops.add(
            ops.divide(ops.multiply(u_t_1, m), 1 - u_product_t_1),
            ops.divide(ops.multiply(1 - u_t, gradient), 1 - u_product_t),
        )
        v_hat = ops.divide(v, (1 - beta_2_power))

        self.assign_sub(
            variable,
            ops.divide(
                ops.multiply(m_hat, lr), ops.add(ops.sqrt(v_hat), 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


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