def METHOD_NAME(self, gradient, variable):
    """Update step given gradient and the associated model variable."""
    lr = tf.cast(self.learning_rate, variable.dtype)
    local_step = tf.cast(self.iterations + 1, variable.dtype)
    beta_1_power = tf.pow(tf.cast(self.beta_1, variable.dtype), local_step)
    var_key = self._var_key(variable)
    m = self._m[self._index_dict[var_key]]
    u = self._u[self._index_dict[var_key]]
    if isinstance(gradient, tf.IndexedSlices):
        # Sparse gradients.
        indices = gradient.indices
        m.assign_add(-m * (1 - self.beta_1))
        m.scatter_add(
            tf.IndexedSlices(gradient.values * (1 - self.beta_1), indices)
        )
        u.assign(u * self.beta_2)
        u_slice = tf.gather(u, indices)
        u_slice_incremental = (
            tf.maximum(u_slice, tf.abs(gradient.values)) - u_slice
        )
        u.scatter_add(tf.IndexedSlices(u_slice_incremental, indices))
        variable.assign_sub(
            (lr * m) / ((1 - beta_1_power) * (u + self.epsilon))
        )
    else:
        # Dense gradients.
        m.assign_add((gradient - m) * (1 - self.beta_1))
        u.assign(tf.maximum(self.beta_2 * u, tf.abs(gradient)))
        variable.assign_sub(
            (lr * m) / ((1 - beta_1_power) * (u + self.epsilon))
        )