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 _shared_step(self):
      with th.no_grad():
          # Frequently alloc and free shared memory to hold intermediate tensor is expensive
          # We cache shared memory buffers in shared_emb.
          shared_emb = {emb.name: ([], []) for emb in self._params}
  
          # Go through all sparse embeddings
          for (
              emb
          ) in self._params:  # pylint: disable=too-many-nested-blocks
              emb_name = emb.name
  
              # we need to combine gradients from multiple forward paths
              idx = []
              grad = []
              for i, data in emb._trace:
                  idx.append(i)
                  grad.append(data.grad.data)
              # If the sparse embedding is not used in the previous forward step
              # The idx and grad will be empty, initialize them as empty tensors to
              # avoid crashing the optimizer step logic.
              #
              # Note: we cannot skip the gradient exchange and update steps as other
              # working processes may send gradient update requests corresponding
              # to certain embedding to this process.
              idx = (
                  th.cat(idx, dim=0)
                  if len(idx) != 0
                  else th.zeros(
                      (0,), dtype=th.long, device=th.device("cpu")
                  )
              )
              grad = (
                  th.cat(grad, dim=0)
                  if len(grad) != 0
                  else th.zeros(
                      (0, emb.embedding_dim),
                      dtype=th.float32,
                      device=th.device("cpu"),
                  )
              )
  
              device = grad.device
              idx_dtype = idx.dtype
              grad_dtype = grad.dtype
              grad_dim = grad.shape[1]
              if self._world_size > 1:
                  if emb_name not in self._shared_cache:
                      self._shared_cache[emb_name] = {}
  
                  # Each training process takes the resposibility of updating a range
                  # of node embeddings, thus we can parallel the gradient update.
                  # The overall progress includes:
                  #   1. In each training process:
                  #     1.a Deciding which process a node embedding belongs to according
                  #         to the formula: process_id = node_idx mod num_of_process(N)
                  #     1.b Split the node index tensor and gradient tensor into N parts
                  #         according to step 1.
                  #     1.c Write each node index sub-tensor and gradient sub-tensor into
                  #         different DGL shared memory buffers.
                  #   2. Cross training process synchronization
                  #   3. In each traning process:
                  #     3.a Collect node index sub-tensors and gradient sub-tensors
                  #     3.b Do gradient update
                  #   4. Done
                  idx_split = th.remainder(idx, self._world_size).long()
                  for i in range(self._world_size):
                      mask = idx_split == i
                      idx_i = idx[mask]
                      grad_i = grad[mask]
  
                      if i == self._rank:
                          shared_emb[emb_name][0].append(idx_i)
                          shared_emb[emb_name][1].append(grad_i)
                      else:
                          # currently nccl does not support Alltoallv operation
                          # we need to use CPU shared memory to share gradient
                          # across processes
                          idx_i = idx_i.to(th.device("cpu"))
                          grad_i = grad_i.to(th.device("cpu"))
                          idx_shmem_name = "idx_{}_{}_{}".format(
                              emb_name, self._rank, i
                          )
                          grad_shmem_name = "grad_{}_{}_{}".format(
                              emb_name, self._rank, i
                          )
  
                          # Create shared memory to hold temporary index and gradient tensor for
                          # cross-process send and recv.
                          if (
                              idx_shmem_name
                              not in self._shared_cache[emb_name]
                              or self._shared_cache[emb_name][
                                  idx_shmem_name
                              ].shape[0]
                              < idx_i.shape[0]
                          ):
                              if (
                                  idx_shmem_name
                                  in self._shared_cache[emb_name]
                              ):
                                  self.shmem_buffer_holder.append(
                                      self._shared_cache[emb_name][
                                          idx_shmem_name
                                      ]
                                  )
                                  self.shmem_buffer_holder.append(
                                      self._shared_cache[emb_name][
                                          grad_shmem_name
                                      ]
                                  )
  
                              # The total number of buffers is the number of NodeEmbeddings *
                              # world_size * (world_size - 1). The minimun buffer size is 128.
                              #
                              # We extend the buffer by idx_i.shape[0] * 2 to avoid
                              # frequent shared memory allocation.
                              # The overall buffer cost will be smaller than three times
                              # the maximum memory requirement for sharing gradients.
                              buffer_size = (
                                  128
                                  if idx_i.shape[0] < 128
                                  else idx_i.shape[0] * 2
                              )
                              idx_shmem = create_shared_mem_array(
                                  "{}_{}".format(
                                      idx_shmem_name, buffer_size
                                  ),
                                  (buffer_size,),
                                  idx_dtype,
                              )
                              grad_shmem = create_shared_mem_array(
                                  "{}_{}".format(
                                      grad_shmem_name, buffer_size
                                  ),
                                  (buffer_size, grad_dim),
                                  grad_dtype,
                              )
                              self._shared_cache[emb_name][
                                  idx_shmem_name
                              ] = idx_shmem
                              self._shared_cache[emb_name][
                                  grad_shmem_name
                              ] = grad_shmem
  
                          # Fill shared memory with temporal index tensor and gradient tensor
                          self._shared_cache[emb_name][idx_shmem_name][
                              : idx_i.shape[0]
                          ] = idx_i
                          self._shared_cache[emb_name][grad_shmem_name][
                              : idx_i.shape[0]
                          ] = grad_i
                          self._opt_meta[emb_name][self._rank][
                              i
                          ] = idx_i.shape[0]
              else:
                  shared_emb[emb_name][0].append(idx)
                  shared_emb[emb_name][1].append(grad)
  
          # make sure the idx shape is passed to each process through opt_meta
          if self._world_size > 1:
              th.distributed.barrier()
          for (
              emb
          ) in self._params:  # pylint: disable=too-many-nested-blocks
              emb_name = emb.name
              if self._world_size > 1:
                  # The first element in shared_emb[emb_name][0] is the local idx
                  device = shared_emb[emb_name][0][0].device
                  # gather gradients from all other processes
                  for i in range(self._world_size):
                      if i != self._rank:
                          idx_shmem_name = "idx_{}_{}_{}".format(
                              emb_name, i, self._rank
                          )
                          grad_shmem_name = "grad_{}_{}_{}".format(
                              emb_name, i, self._rank
                          )
                          size = self._opt_meta[emb_name][i][self._rank]
  
                          # Retrive shared memory holding the temporal index and gradient
                          # tensor that is sent to current training process
                          if (
                              idx_shmem_name
                              not in self._shared_cache[emb_name]
                              or self._shared_cache[emb_name][
                                  idx_shmem_name
                              ].shape[0]
                              < size
                          ):
                              buffer_size = (
                                  128 if size < 128 else size * 2
                              )
                              idx_shmem = get_shared_mem_array(
                                  "{}_{}".format(
                                      idx_shmem_name, buffer_size
                                  ),
                                  (buffer_size,),
                                  idx_dtype,
                              )
                              grad_shmem = get_shared_mem_array(
                                  "{}_{}".format(
                                      grad_shmem_name, buffer_size
                                  ),
                                  (buffer_size, grad_dim),
                                  grad_dtype,
                              )
                              self._shared_cache[emb_name][
                                  idx_shmem_name
                              ] = idx_shmem
                              self._shared_cache[emb_name][
                                  grad_shmem_name
                              ] = grad_shmem
  
                          idx_i = self._shared_cache[emb_name][
                              idx_shmem_name
                          ][:size]
                          grad_i = self._shared_cache[emb_name][
                              grad_shmem_name
                          ][:size]
                          shared_emb[emb_name][0].append(
                              idx_i.to(device, non_blocking=True)
                          )
                          shared_emb[emb_name][1].append(
                              grad_i.to(device, non_blocking=True)
                          )
  
          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 = th.cat(shared_emb[emb_name][0], dim=0)
              grad = th.cat(shared_emb[emb_name][1], dim=0)
              self.update(idx, grad, emb)
  
          # synchronized gradient update
          if self._world_size > 1:
              th.distributed.barrier()
  
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 7 for calling the SparseGradOptimizer obj's underlying member method code with example usage: obj._shared_step() without return types