code: def _all_gather(tensor: torch.Tensor, comm_spec: CommSpec):
      """
      Implement all gather operation on device mesh based on information provided by comm_spec.
      """
      process_group = comm_spec.process_group_dict[
          comm_spec.logical_process_axis
      ]
      world_size = dist.get_world_size(process_group)
      tensor_list = [
          torch.zeros(
              tensor.shape, dtype=tensor.dtype, device=tensor.device
          )
          for _ in range(world_size)
      ]
      # without this contiguous operation, the all gather may get some unexpected results.
      tensor = tensor.contiguous()
      dist.all_gather(tensor_list, tensor, group=process_group)
      output = torch.cat(
          tuple(tensor_list), comm_spec.gather_dim
      ).contiguous()
      return output
  
imports: from enum import Enum
  from typing import Dict
  import torch
  import torch.distributed
  from torch.distributed import ReduceOp
label_desc: Use image node_id 1 for calling a global function with example usage: _all_gather(tensor, comm_spec) and returns: output