code: def update(self, index, weight, grad, state):
      assert isinstance(weight, NDArray)
      assert isinstance(grad, NDArray)
      self._update_count(index)
      lr = self._get_lr(index)
      wd = self._get_wd(index)
  
      t = self._index_update_count[index]
  
      # preprocess grad
      # grad = grad * self.rescale_grad + wd * weight
      grad *= self.rescale_grad + wd * weight
      if self.clip_gradient is not None:
          grad = clip(grad, -self.clip_gradient, self.clip_gradient)
  
      # warming momentum schedule
      momentum_t = self.beta1 * (
          1.0 - 0.5 * (pow(0.96, t * self.schedule_decay))
      )
      momentum_t_1 = self.beta1 * (
          1.0 - 0.5 * (pow(0.96, (t + 1) * self.schedule_decay))
      )
      self.m_schedule = self.m_schedule * momentum_t
      m_schedule_next = self.m_schedule * momentum_t_1
  
      # update m_t and v_t
      m_t, v_t = state
      m_t[:] = self.beta1 * m_t + (1.0 - self.beta1) * grad
      v_t[:] = self.beta2 * v_t + (1.0 - self.beta2) * grad * grad
  
      grad_prime = grad / (1.0 - self.m_schedule)
      m_t_prime = m_t / (1.0 - m_schedule_next)
      v_t_prime = v_t / (1.0 - pow(self.beta2, t))
      m_t_bar = (
          1.0 - momentum_t
      ) * grad_prime + momentum_t_1 * m_t_prime
  
      # update weight
      weight[:] -= lr * m_t_bar / (sqrt(v_t_prime) + self.epsilon)
  
file_path: insightface/alignment/heatmap/optimizer.py
imports: import mxnet
  import mxnet.optimizer
  from mxnet.ndarray import NDArray, zeros, clip, sqrt, cast, maximum, abs
label_desc: Use image node_id 3 for calling the ONadam obj's underlying member method code with example usage: obj.update(index, weight, grad, state) without return types