def _all_get_tensor(self, shared_name, tensor, shape):
        """A helper function to get model-parallel tensors.

        This method must and only need to be called in multi-GPU DDP training.
        For now, it's only used in ``all_get_embedding`` and
        ``_all_get_optm_state``.
        """
        # create a shared memory tensor
        if self._rank == 0:
            # root process creates shared memory
            val = create_shared_mem_array(
                shared_name,
                shape,
                tensor.dtype,
            )
            self._store.set(shared_name, shared_name)
        else:
            self._store.wait([shared_name])
            val = get_shared_mem_array(
                shared_name,
                shape,
                tensor.dtype,
            )
        # need to map indices and slice into existing tensor
        idxs = self._partition.map_to_global(
            F.arange(0, tensor.shape[0], ctx=F.context(tensor)),
            self._rank,
        ).to(val.device)
        val[idxs] = tensor.to(val.device)

        self._store.delete_key(shared_name)
        # wait for all processes to finish
        th.distributed.barrier()
        return val
A helper function to get model-parallel tensors.

        This method must and only need to be called in multi-GPU DDP training.
        For now, it's only used in ``all_get_embedding`` and
        ``_all_get_optm_state``.
        
python/dgl/nn/pytorch/sparse_emb.py