class_docstr: The abstract sparse optimizer.
  
  Note: dgl sparse optimizer only work with dgl.NodeEmbedding
  
  Parameters
  ----------
  params : list of NodeEmbedding
      The list of NodeEmbeddings.
  lr : float
      The learning rate.
code: def _comm_step(self):
      with th.no_grad():
          idx_in = {}
          grad_in = {}
          for (
              emb
          ) in self._params:  # pylint: disable=too-many-nested-blocks
              emb_name = emb.name
              partition = emb.partition
  
              if not partition:
                  # use default partitioning
                  partition = NDArrayPartition(
                      emb.num_embeddings,
                      self._world_size if self._world_size > 0 else 1,
                      mode="remainder",
                  )
  
              # we need to combine gradients from multiple forward paths
              if len(emb._trace) == 0:
                  idx = th.zeros(
                      (0,), dtype=th.long, device=self._device
                  )
                  grad = th.zeros(
                      (0, emb.embedding_dim),
                      dtype=th.float32,
                      device=self._device,
                  )
              elif len(emb._trace) == 1:
                  # the special case where we can use the tensors as is
                  # without any memcpy's
                  idx, grad = emb._trace[0]
                  grad = grad.grad.data
              else:
                  idx = []
                  grad = []
                  for i, data in emb._trace:
                      idx.append(i)
                      grad.append(data.grad.data)
                  idx = th.cat(idx, dim=0)
                  grad = th.cat(grad, dim=0)
  
              (
                  idx_in[emb_name],
                  grad_in[emb_name],
              ) = nccl.sparse_all_to_all_push(
                  idx, grad, partition=partition
              )
              if emb.partition:
                  # if the embedding is partitioned, map back to indexes
                  # into the local tensor
                  idx_in[emb_name] = partition.map_to_local(
                      idx_in[emb_name]
                  )
  
          if self._clean_grad:
              # clean gradient track
              for emb in self._params:
                  emb.reset_trace()
              self._clean_grad = False
  
          for emb in self._params:
              emb_name = emb.name
              idx = idx_in[emb_name]
              grad = grad_in[emb_name]
              self.update(idx, grad, emb)
  
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 6 for calling the SparseGradOptimizer obj's underlying member method code with example usage: obj._comm_step() without return types