class_docstr: Class for storing node embeddings.
  
  The class is optimized for training large-scale node embeddings. It updates the embedding in
  a sparse way and can scale to graphs with millions of nodes. It also supports partitioning
  to multiple GPUs (on a single machine) for more acceleration. It does not support partitioning
  across machines.
  
  Currently, DGL provides two optimizers that work with this NodeEmbedding
  class: ``SparseAdagrad`` and ``SparseAdam``.
  
  The implementation is based on torch.distributed package. It depends on the pytorch
  default distributed process group to collect multi-process information and uses
  ``torch.distributed.TCPStore`` to share meta-data information across multiple gpu processes.
  It use the local address of '127.0.0.1:12346' to initialize the TCPStore.
  
  NOTE: The support of NodeEmbedding is experimental.
  
  Parameters
  ----------
  num_embeddings : int
      The number of embeddings. Currently, the number of embeddings has to be the same as
      the number of nodes.
  embedding_dim : int
      The dimension size of embeddings.
  name : str
      The name of the embeddings. The name should uniquely identify the embeddings in the system.
  init_func : callable, optional
      The function to create the initial data. If the init function is not provided,
      the values of the embeddings are initialized to zero.
  device : th.device
      Device to store the embeddings on.
  parittion : NDArrayPartition
      The partition to use to distributed the embeddings between
      processes.
  
  Examples
  --------
  Before launching multiple gpu processes
  
  >>> def initializer(emb):
          th.nn.init.xavier_uniform_(emb)
          return emb
  
  In each training process
  
  >>> 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()
class_docstr_tok: ["Class","for","storing","node","embeddings",".","The","class","is","optimized","for","training","large-scale","node","embeddings",".","It","updates","the","embedding","in","a","sparse","way","and","can","scale","to","graphs","with","millions","of","nodes",".","It","also","supports","partitioning","to","multiple","GPUs","(","on","a","single","machine",")","for","more","acceleration",".","It","does","not","support","partitioning","across","machines",".","Currently",",","DGL","provides","two","optimizers","that","work","with","this","NodeEmbedding","class",":","``","SparseAdagrad","``","and","``","SparseAdam","``",".","The","implementation","is","based","on","torch.distributed","package",".","It","depends","on","the","pytorch","default","distributed","process","group","to","collect","multi-process","information","and","uses","``","torch.distributed.TCPStore","``","to","share","meta-data","information","across","multiple","gpu","processes",".","It","use","the","local","address","of","'127.0.0.1:12346","'","to","initialize","the","TCPStore",".","NOTE",":","The","support","of","NodeEmbedding","is","experimental",".","Parameters","--","--","--","--","--","num_embeddings",":","int","The","number","of","embeddings",".","Currently",",","the","number","of","embeddings","has","to","be","the","same","as","the","number","of","nodes",".","embedding_dim",":","int","The","dimension","size","of","embeddings",".","name",":","str","The","name","of","the","embeddings",".","The","name","should","uniquely","identify","the","embeddings","in","the","system",".","init_func",":","callable",",","optional","The","function","to","create","the","initial","data",".","If","the","init","function","is","not","provided",",","the","values","of","the","embeddings","are","initialized","to","zero",".","device",":","th.device","Device","to","store","the","embeddings","on",".","parittion",":","NDArrayPartition","The","partition","to","use","to","distributed","the","embeddings","between","processes",".","Examples","--","--","--","--","Before","launching","multiple","gpu","processes",">",">",">","def","initializer","(","emb",")",":","th.nn.init.xavier_uniform_","(","emb",")","return","emb","In","each","training","process",">",">",">","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","(",")"]
code: def __init__(
      self,
      num_embeddings,
      embedding_dim,
      name,
      init_func=None,
      device=None,
      partition=None,
  ):
      global _STORE
  
      if device is None:
          device = th.device("cpu")
  
      # Check whether it is multi-gpu training or not.
      if th.distributed.is_initialized():
          rank = th.distributed.get_rank()
          world_size = th.distributed.get_world_size()
      else:
          rank = -1
          world_size = 0
      self._rank = rank
      self._world_size = world_size
      self._store = None
      self._comm = None
      self._partition = partition
  
      host_name = "127.0.0.1"
      port = 12346
  
      if rank >= 0:
          # for multi-gpu training, setup a TCPStore for
          # embeding status synchronization across GPU processes
          if _STORE is None:
              _STORE = th.distributed.TCPStore(
                  host_name,
                  port,
                  world_size,
                  rank == 0,
                  timedelta(seconds=10 * 60),
              )
          self._store = _STORE
  
      # embeddings is stored in CPU memory.
      if th.device(device) == th.device("cpu"):
          if rank <= 0:
              emb = create_shared_mem_array(
                  name, (num_embeddings, embedding_dim), th.float32
              )
              if init_func is not None:
                  emb = init_func(emb)
          if rank == 0:  # the master gpu process
              for _ in range(1, world_size):
                  # send embs
                  self._store.set(name, name)
          elif rank > 0:
              # receive
              self._store.wait([name])
              emb = get_shared_mem_array(
                  name, (num_embeddings, embedding_dim), th.float32
              )
          self._tensor = emb
      else:  # embeddings is stored in GPU memory.
          self._comm = True
  
          if not self._partition:
              # for communication we need a partition
              self._partition = NDArrayPartition(
                  num_embeddings,
                  self._world_size if self._world_size > 0 else 1,
                  mode="remainder",
              )
  
          # create local tensors for the weights
          local_size = self._partition.local_size(max(self._rank, 0))
  
          # TODO(dlasalle): support 16-bit/half embeddings
          emb = th.empty(
              [local_size, embedding_dim],
              dtype=th.float32,
              requires_grad=False,
              device=device,
          )
          if init_func:
              emb = init_func(emb)
          self._tensor = emb
  
      self._num_embeddings = num_embeddings
      self._embedding_dim = embedding_dim
      self._name = name
      self._optm_state = None  # track optimizer state
      self._trace = []
  
imports: from datetime import timedelta
  import torch
  from ...backend import pytorch
  from ...cuda import nccl
  from ...partition import NDArrayPartition
  from ...utils import create_shared_mem_array, get_shared_mem_array
label_desc: Use image node_id 1 to create a new NodeEmbedding object  with example: obj = NodeEmbedding(num_embeddings, embedding_dim, name, init_func, device, partition)