def METHOD_NAME(param, moment, grad_sq_sum, grad, lr, momentum, beta):
    grad_sq_sum_out = grad_sq_sum + np.sum(grad * grad)
    nlr = lr * np.power(beta + grad_sq_sum_out, -1.0 / 3.0)
    alpha = momentum * np.square(nlr)
    moment_out = grad + (1 - alpha) * (moment - grad)
    param_out = param + nlr * moment_out
    return (param_out.astype(np.float32), moment_out.astype(np.float32),
            grad_sq_sum_out.astype(np.float32))