'Update the parameters.
Parameters
index : int
An unique integer key used to index the parameters
weight : NDArray
weight ndarray
grad : NDArray
grad ndarray
state : NDArray or other objects returned by init_state
The auxiliary state used in optimization.'

def update(self, index, weight, grad, state):

    lr = self._get_lr(index)
    self._update_count(index)
    t = self._index_update_count[index]
    (mean, variance) = state
    wd = self._get_wd(index)
    grad = ((grad * self.rescale_grad) + (wd * weight))
    if (self.clip_gradient is not None):
        mx.nd.clip(grad, (- self.clip_gradient), self.clip_gradient, out=grad)
    mean *= self.beta1
    mean += (grad * (1.0 - self.beta1))
    variance *= self.beta2
    variance += ((1 - self.beta2) * mx.nd.square(grad, out=grad))
    coef1 = (1.0 - (self.beta1 ** t))
    coef2 = (1.0 - (self.beta2 ** t))
    lr *= (math.sqrt(coef2) / coef1)
    weight -= ((lr * mean) / (mx.nd.sqrt(variance) + self.epsilon))
