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
code: def __init__(self, params, lr, betas=(0.9, 0.999), eps=1e-08):
      super(SparseAdam, self).__init__(params, lr)
      self._eps = eps
      # We need to register a state sum for each embedding in the kvstore.
      self._beta1 = betas[0]
      self._beta2 = betas[1]
      self._defaults = {
          "_lr": lr,
          "_eps": eps,
          "_beta1": betas[0],
          "_beta2": betas[1],
      }
      for emb in params:
          assert isinstance(
              emb, DistEmbedding
          ), "SparseAdam only supports dgl.distributed.DistEmbedding"
  
          state_step = DistTensor(
              (emb.num_embeddings,),
              th.float32,
              emb.name + "_step",
              init_func=initializer,
              part_policy=emb.part_policy,
              is_gdata=False,
          )
          state_mem = DistTensor(
              (emb.num_embeddings, emb.embedding_dim),
              th.float32,
              emb.name + "_mem",
              init_func=initializer,
              part_policy=emb.part_policy,
              is_gdata=False,
          )
          state_power = DistTensor(
              (emb.num_embeddings, emb.embedding_dim),
              th.float32,
              emb.name + "_power",
              init_func=initializer,
              part_policy=emb.part_policy,
              is_gdata=False,
          )
          state = (state_step, state_mem, state_power)
          assert (
              emb.name not in self._state
          ), "{} already registered in the optimizer".format(emb.name)
          self._state[emb.name] = state
  
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 1 to create a new SparseAdam object from inherited base classes: DistSparseGradOptimizer with example: obj = SparseAdam(params, lr, betas, eps)