def METHOD_NAME(self, gradient, variable, learning_rate):
    """Update step given gradient and the associated model variable."""
    lr = ops.cast(learning_rate, variable.dtype)
    gradient = ops.cast(gradient, variable.dtype)
    epsilon_2 = ops.cast(self.epsilon_2, variable.dtype)
    one = ops.cast(1.0, variable.dtype)
    local_step = ops.cast(self.iterations + 1, variable.dtype)
    if not callable(self._learning_rate) and self.relative_step:
        lr = ops.minimum(lr, 1 / ops.sqrt(local_step))
    r = self._r[self._get_variable_index(variable)]
    c = self._c[self._get_variable_index(variable)]
    v = self._v[self._get_variable_index(variable)]
    rho_t = ops.minimum(lr, 1 / ops.sqrt(local_step))
    alpha_t = ops.maximum(epsilon_2, self._rms(variable)) * rho_t
    regulated_grad_square = ops.square(gradient) + self.epsilon_1
    beta_2_t = 1 - ops.power(local_step, self.beta_2_decay)
    if len(variable.shape) >= 2:
        # `r` deletes the last dimension of gradient, so it is of shape
        # `gradient.shape[:-1]`.
        r.assign(
            beta_2_t * r
            + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-1)
        )
        # `c` deletes the second last dimension of gradient, so it is of
        # shape `gradient.shape[:-2] + gradient.shape[-1]`.
        c.assign(
            beta_2_t * c
            + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-2)
        )
        v.assign(
            ops.expand_dims(
                r / ops.mean(r, axis=-1, keepdims=True), axis=-1
            )
            * ops.expand_dims(c, -2)
        )
    else:
        v.assign(beta_2_t * v + (1 - beta_2_t) * regulated_grad_square)
    # `convert_to_tensor` unifies the handling of sparse and dense grads.
    u_t = gradient / ops.sqrt(v)
    u_t_hat = u_t / ops.maximum(one, (self._rms(u_t) / self.clip_threshold))
    variable.assign(variable - alpha_t * u_t_hat)