class_docstr: Distributed Node embedding optimizer using the Adam algorithm.
  
  This optimizer implements a distributed sparse version of Adam algorithm for
  optimizing :class:`dgl.distributed.DistEmbedding`. 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.distributed.DistEmbedding]
      The list of dgl.distributed.DistEmbedding.
  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
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.distributed.DistEmbedding
      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.distributed.DistEmbedding","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.distributed.DistEmbedding
          Sparse embedding to update.
      """
      beta1 = self._beta1
      beta2 = self._beta2
      eps = self._eps
      clr = self._lr
      state_step, state_mem, state_power = self._state[emb.name]
  
      state_dev = th.device("cpu")
      exec_dev = grad.device
  
      # 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 = (
          state_dev == th.device("cpu") and exec_dev != state_dev
      )
  
      # the update is non-linear so indices must be unique
      grad_indices, inverse, cnt = th.unique(
          idx, return_inverse=True, return_counts=True
      )
      # update grad state
      state_idx = grad_indices.to(state_dev)
      # The original implementation will cause read/write contension.
      #    state_step[state_idx] += 1
      #    state_step = state_step[state_idx].to(exec_dev, non_blocking=True)
      # In a distributed environment, the first line of code will send write requests to
      # kvstore servers to update the state_step which is asynchronous and the second line
      # of code will also send read requests to kvstore servers. The write and read requests
      # may be handled by different kvstore servers managing the same portion of the
      # state_step dist tensor in the same node. So that, the read request may read an old
      # value (i.e., 0 in the first iteration) which will cause
      # update_power_corr to be NaN
      state_val = state_step[state_idx] + 1
      state_step[state_idx] = state_val
      state_step = state_val.to(exec_dev)
      orig_mem = state_mem[state_idx].to(exec_dev)
      orig_power = state_power[state_idx].to(exec_dev)
  
      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
      update_mem_dst = update_mem.to(state_dev, non_blocking=True)
      update_power_dst = update_power.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()
          # 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._tensor[state_idx] -= std_values_dst
  
imports: import abc
  import warnings
  from abc import abstractmethod
  from os.path import exists
  import torch
  import dgl
  from ....None import backend
  from ...dist_tensor import DistTensor
  from ...graph_partition_book import EDGE_PART_POLICY, NODE_PART_POLICY
  from ...nn.pytorch import DistEmbedding
  from .utils import alltoall, alltoallv
label_desc: Use image node_id 2 for calling the SparseAdam obj's underlying member method code with example usage: obj.update(idx, grad, emb) without return types