class_docstr: Node embedding optimizer using the Adam algorithm.
  
  This optimizer implements a sparse version of Adagrad algorithm for
  optimizing :class:`dgl.nn.NodeEmbedding`. Being sparse means it only
  updates the embeddings whose gradients have updates, which are usually
  a very small portion of the total embeddings.
  
  Adam maintains a :math:`Gm_{t,i,j}` and `Gp_{t,i,j}` for every parameter
  in the embeddings, where
  :math:`Gm_{t,i,j}=beta1 * Gm_{t-1,i,j} + (1-beta1) * g_{t,i,j}`,
  :math:`Gp_{t,i,j}=beta2 * Gp_{t-1,i,j} + (1-beta2) * g_{t,i,j}^2`,
  :math:`g_{t,i,j} = lr * Gm_{t,i,j} / (1 - beta1^t) / \sqrt{Gp_{t,i,j} / (1 - beta2^t)}` and
  :math:`g_{t,i,j}` is the gradient of the dimension :math:`j` of embedding :math:`i`
  at step :math:`t`.
  
  NOTE: The support of sparse Adam optimizer is experimental.
  
  Parameters
  ----------
  params : list[dgl.nn.NodeEmbedding]
      The list of dgl.nn.NodeEmbeddings.
  lr : float
      The learning rate.
  betas : tuple[float, float], Optional
      Coefficients used for computing running averages of gradient and its square.
      Default: (0.9, 0.999)
  eps : float, Optional
      The term added to the denominator to improve numerical stability
      Default: 1e-8
  use_uva : bool, Optional
      Whether to use pinned memory for storing 'mem' and 'power' parameters,
      when the embedding is stored on the CPU. This will improve training
      speed, but will require locking a large number of virtual memory pages.
      For embeddings which are stored in GPU memory, this setting will have
      no effect.
      Default: True if the gradients are generated on the GPU, and False
      if the gradients are on the CPU.
  dtype : torch.dtype, Optional
      The type to store optimizer state with. Default: th.float32.
  
  Examples
  --------
  >>> def initializer(emb):
          th.nn.init.xavier_uniform_(emb)
          return emb
  >>> emb = dgl.nn.NodeEmbedding(g.num_nodes(), 10, 'emb', init_func=initializer)
  >>> optimizer = dgl.optim.SparseAdam([emb], lr=0.001)
  >>> for blocks in dataloader:
  ...     ...
  ...     feats = emb(nids, gpu_0)
  ...     loss = F.sum(feats + 1, 0)
  ...     loss.backward()
  ...     optimizer.step()
docstr: Update embeddings in a sparse manner
  Sparse embeddings are updated in mini batches. We maintain gradient states for
  each embedding so they can be updated separately.
  
  Parameters
  ----------
  idx : tensor
      Index of the embeddings to be updated.
  grad : tensor
      Gradient of each embedding.
  emb : dgl.nn.NodeEmbedding
      Sparse embedding to update.
docstr_tok: ["Update","embeddings","in","a","sparse","manner","Sparse","embeddings","are","updated","in","mini","batches",".","We","maintain","gradient","states","for","each","embedding","so","they","can","be","updated","separately",".","Parameters","--","--","--","--","--","idx",":","tensor","Index","of","the","embeddings","to","be","updated",".","grad",":","tensor","Gradient","of","each","embedding",".","emb",":","dgl.nn.NodeEmbedding","Sparse","embedding","to","update","."]
code: def update(self, idx, grad, emb):
      """Update embeddings in a sparse manner
      Sparse embeddings are updated in mini batches. We maintain gradient states for
      each embedding so they can be updated separately.
  
      Parameters
      ----------
      idx : tensor
          Index of the embeddings to be updated.
      grad : tensor
          Gradient of each embedding.
      emb : dgl.nn.NodeEmbedding
          Sparse embedding to update.
      """
      with th.no_grad():
          state_step, state_mem, state_power = emb.optm_state
          exec_dtype = grad.dtype
          exec_dev = grad.device
          state_dev = state_step.device
  
          # whether or not we need to transfer data from the GPU to the CPU
          # while updating the weights
          is_d2h = state_dev.type == "cpu" and exec_dev.type == "cuda"
  
          # only perform async copies cpu -> gpu, or gpu-> gpu, but block
          # when copying to the cpu, so as to ensure the copy is finished
          # before operating on the data on the cpu
          state_block = is_d2h
  
          if self._is_using_uva[emb.name] is None and is_d2h:
              # we should use UVA going forward
              self._setup_uva(emb.name, state_mem, state_power)
          elif self._is_using_uva[emb.name] is None:
              # we shouldn't use UVA going forward
              self._is_using_uva[emb.name] = False
  
          use_uva = self._is_using_uva[emb.name]
  
          beta1 = self._beta1
          beta2 = self._beta2
          eps = self._eps
  
          clr = self._lr
          # There can be duplicated indices due to sampling.
          # Thus unique them here and average the gradient here.
          grad_indices, inverse, cnt = th.unique(
              idx, return_inverse=True, return_counts=True
          )
          state_idx = grad_indices.to(state_dev)
          state_step[state_idx] += 1
          state_step = state_step[state_idx].to(exec_dev)
  
          if use_uva:
              orig_mem = gather_pinned_tensor_rows(
                  state_mem, grad_indices
              )
              orig_power = gather_pinned_tensor_rows(
                  state_power, grad_indices
              )
          else:
              orig_mem = state_mem[state_idx].to(exec_dev)
              orig_power = state_power[state_idx].to(exec_dev)
          # convert to exec dtype
          orig_mem = orig_mem.to(dtype=exec_dtype)
          orig_power = orig_power.to(dtype=exec_dtype)
  
          grad_values = th.zeros(
              (grad_indices.shape[0], grad.shape[1]), device=exec_dev
          )
          grad_values.index_add_(0, inverse, grad)
          grad_values = grad_values / cnt.unsqueeze(1)
  
          grad_mem = grad_values
          grad_power = grad_values * grad_values
  
          update_mem = beta1 * orig_mem + (1.0 - beta1) * grad_mem
          update_power = beta2 * orig_power + (1.0 - beta2) * grad_power
  
          if use_uva:
              scatter_pinned_tensor_rows(
                  state_mem,
                  grad_indices,
                  update_mem.to(dtype=self._dtype),
              )
              scatter_pinned_tensor_rows(
                  state_power,
                  grad_indices,
                  update_power.to(dtype=self._dtype),
              )
          else:
              update_mem_dst = update_mem.to(dtype=self._dtype).to(
                  state_dev, non_blocking=True
              )
              update_power_dst = update_power.to(dtype=self._dtype).to(
                  state_dev, non_blocking=True
              )
              if state_block:
                  # use events to try and overlap CPU and GPU as much as possible
                  update_event = th.cuda.Event()
                  update_event.record()
  
          update_mem_corr = update_mem / (
              1.0
              - th.pow(th.tensor(beta1, device=exec_dev), state_step)
          ).unsqueeze(1)
          update_power_corr = update_power / (
              1.0
              - th.pow(th.tensor(beta2, device=exec_dev), state_step)
          ).unsqueeze(1)
          std_values = (
              clr * update_mem_corr / (th.sqrt(update_power_corr) + eps)
          )
          std_values_dst = std_values.to(state_dev, non_blocking=True)
  
          if state_block:
              std_event = th.cuda.Event()
              std_event.record()
  
          if not use_uva:
              if state_block:
                  # wait for our transfers from exec_dev to state_dev to finish
                  # before we can use them
                  update_event.wait()
              state_mem[state_idx] = update_mem_dst
              state_power[state_idx] = update_power_dst
  
          if state_block:
              # wait for the transfer of std_values to finish before we
              # can use it
              std_event.wait()
          emb.weight[state_idx] -= std_values_dst
  
imports: import abc
  from abc import abstractmethod
  import torch
  from ...cuda import nccl
  from ...nn.pytorch import NodeEmbedding
  from ...partition import NDArrayPartition
  from ...utils import create_shared_mem_array, gather_pinned_tensor_rows, get_shared_mem_array, pin_memory_inplace, scatter_pinned_tensor_rows
label_desc: Use image node_id 4 for calling the SparseAdam obj's underlying member method code with example usage: obj.update(idx, grad, emb) without return types