docstr: _gather gather sharded tensors to a replicated one.
  Args:
      tensor (torch.Tensor): a shared torch tensor
      old_dist_spec (_DistSpec): the distributed spec. of the tensor.
  
  Returns:
      torch.Tensor: a replicated tensor.
docstr_tok: ["_gather","gather","sharded","tensors","to","a","replicated","one",".","Args",":","tensor","(","torch.Tensor",")",":","a","shared","torch","tensor","old_dist_spec","(","_DistSpec",")",":","the","distributed","spec",".","of","the","tensor",".","Returns",":","torch.Tensor",":","a","replicated","tensor","."]
code: def _gather(
      tensor: torch.Tensor, old_dist_spec: _DistSpec, pg: ProcessGroup
  ) -> torch.Tensor:
      """_gather gather sharded tensors to a replicated one.
      Args:
          tensor (torch.Tensor): a shared torch tensor
          old_dist_spec (_DistSpec): the distributed spec. of the tensor.
  
      Returns:
          torch.Tensor: a replicated tensor.
      """
      assert (
          old_dist_spec.placement.value == "s"
      ), f"The old_dist_spec of DistSpecManager._gather must be SHARD!"
      is_cpu_tensor = False
      if tensor.device.type == "cpu":
          # pytorch lower than 1.11 dose not support gather a cpu tensor.
          # Therefore, we transfer tensor to GPU before gather.
          saved_dev = tensor.device
          tensor.data = tensor.data.cuda()
          is_cpu_tensor = True
  
      buffer = [
          torch.empty_like(tensor) for _ in range(pg.tp_world_size())
      ]
      assert tensor.device.type == "cuda"
      dist.all_gather(buffer, tensor, group=pg.tp_process_group())
      for i in range(len(old_dist_spec.dims) - 1, -1, -1):
          new_buffer = []
          dim = old_dist_spec.dims[i]
          num_parts = old_dist_spec.num_partitions[i]
          for start in range(0, len(buffer), num_parts):
              new_buffer.append(
                  torch.cat(buffer[start : start + num_parts], dim)
              )
          buffer = new_buffer
      assert len(buffer) == 1
  
      if is_cpu_tensor:
          buffer[0].data = buffer[0].data.to(saved_dev)
      return buffer[0]
  
imports: from contextlib import contextmanager
  import torch
  import torch.distributed
  from numpy import prod
  from colossalai.legacy.tensor.distspec import DistPlacementPattern, _DistSpec
  from colossalai.legacy.tensor.process_group import ProcessGroup
label_desc: Use image node_id 3 for calling the DistSpecManager obj's underlying member method code with example usage: obj._gather(tensor, old_dist_spec, pg) and returns: buffer