def update_step(self, gradient, variable, learning_rate):
        """Update step given gradient and the associated model variable."""
        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)
            ),
        )
Update step given gradient and the associated model variable.
keras/src/optimizers/nadam.py