def METHOD_NAME(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)]
    m.assign(m + (gradient - m) * (1 - beta_1))
    v.assign(v + (ops.square(gradient) - v) * (1 - beta_2))
    m_hat = u_t_1 * m / (1 - u_product_t_1) + (1 - u_t) * gradient / (
        1 - u_product_t
    )
    v_hat = v / (1 - beta_2_power)
    variable.assign(
        variable - (m_hat * lr) / (ops.sqrt(v_hat) + self.epsilon)
    )