code: def _mix_gather(tensor, comm_spec):
      """
      Implement mix gather operation on device mesh based on information provided by comm_spec.
      Mix gather is the all-gather operation on all devices in the device_mesh(FlattenDeviceMesh) of the comm_spec. It is
      different from _all_gather because _mix_gather does all-gather in two dimensions of device mesh, while _all_gather
      only does all-gather in one dimension.
      Assume index of f and b target pairs are 'f' and 'b'
      ShardingSpec => gather_dim, logical_process_axes
      S0S1 => [b, f], (1, 0)
      S1S0 => [b, f], (0, 1)
      S01R => [f], (1, 1)
      RS01 => [b], (1, 1)
      Example:
      mesh_shape = (2,4)
              # [[0, 1, 2, 3],
              #  [4, 5, 6, 7]]
              # return {0: [0, 4, 1, 5, 2, 6, 3, 7], 1: [0, 1, 2, 3, 4, 5, 6, 7]}
      S0S1:
      leading_group_dim = 1
      process_group = "[0, 1, 2, 3, 4, 5, 6, 7]"
      tensor_list = [(0,0),(0,1),(0,2),(0,3),(1,0),(1,1),(1,2),(1,3)] # [(slice_id_f, slice_id_b),...]
      mesh_shape = (2,4)
      cat_slice = [4,2]
      tmp_tensor_list = [(...,shape[f],shape[b]*4,...),(...,shape[f],shape[b]*4,...)]
      tmp_tensor_list[0] = torch.cat(((0,0),(0,1),(0,2),(0,3)), dim=b)
      tmp_tensor_list[1] = torch.cat(((1,0),(1,1),(1,2),(1,3)), dim=b)
      output = torch.cat((tmp_tensor_list[0],tmp_tensor_list[1]), dim=a)
      S1S0:
      leading_group_dim = 0
      process_group = "[0, 4, 1, 5, 2, 6, 3, 7]"
      tensor_list = [(0,0),(0,1),(1,0),(1,1),(2,0),(2,1),(3,0),(3,1)]
      mesh_shape = (2,4)
      cat_slice = [2,4]
      tmp_tensor_list = [(...,shape[f],shape[b]*2,...),(...,shape[f],shape[b]*2,...),(...,shape[f],shape[b]*2,...),(...,shape[f],shape[b]*2,...)]
      tmp_tensor_list[0] = torch.cat(((0,0),(0,1)), dim=b)
      tmp_tensor_list[1] = torch.cat(((1,0),(1,1)), dim=b)
      tmp_tensor_list[2] = torch.cat(((2,0),(2,1)), dim=b)
      tmp_tensor_list[3] = torch.cat(((3,0),(3,1)), dim=b)
      S10R:
      leading_group_dim = 0
      process_group = "[0, 4, 1, 5, 2, 6, 3, 7]"
      tensor_list = [(0,0),(1,0),(2,0),(3,0),(4,0),(5,0),(6,0),(7,0)]
      S01R:
      leading_group_dim = 1
      process_group = "[0, 1, 2, 3, 4, 5, 6, 7]"
      tensor_list = [(0,0),(1,0),(2,0),(3,0),(4,0),(5,0),(6,0),(7,0)]
      """
      total_slices = comm_spec.device_mesh.shape[0]
      tensor_list = [
          torch.zeros(
              tensor.shape, dtype=tensor.dtype, device=tensor.device
          )
          for _ in range(total_slices)
      ]
      leading_group_dim = comm_spec.logical_process_axes[0]
      assert len(comm_spec.device_mesh.process_groups_dict) == 1
      _, process_group = comm_spec.device_mesh.process_groups_dict[0][0]
      process_number_list = comm_spec.device_meshes.process_number_dict[
          leading_group_dim
      ]
  
      # Global all_gather
      dist.all_gather(tensor_list, tensor, group=process_group)
  
      # This is very ugly. I'm figuring out more elegant methods
      tensor_list_sorted = [
          torch.zeros(
              tensor.shape, dtype=tensor.dtype, device=tensor.device
          )
          for _ in range(total_slices)
      ]
      for i in range(total_slices):
          tensor_list_sorted[i] = tensor_list[process_number_list[i]]
      tensor_list = tensor_list_sorted
  
      if (
          comm_spec.logical_process_axes[0]
          == comm_spec.logical_process_axes[1]
      ):
          output = torch.cat(
              tuple(tensor_list), comm_spec.gather_dim[0]
          ).contiguous()
      else:
          mesh_shape = comm_spec.device_meshes.shape
          cat_slice = [
              mesh_shape[comm_spec.logical_process_axes[0]],
              mesh_shape[comm_spec.logical_process_axes[1]],
          ]
          tmp_tensor_shape = list(tensor.shape)
          tmp_tensor_shape[comm_spec.gather_dim[0]] *= cat_slice[0]
          tmp_tensor_shape = torch.Size(tmp_tensor_shape)
          tmp_tensor_list = [
              torch.zeros(
                  tmp_tensor_shape,
                  dtype=tensor.dtype,
                  device=tensor.device,
              )
              for _ in range(cat_slice[1])
          ]
          for i in range(cat_slice[1]):
              tmp_tensor_list[i] = torch.cat(
                  tuple(
                      tensor_list[
                          i * cat_slice[0] : (i + 1) * cat_slice[0]
                      ]
                  ),
                  comm_spec.gather_dim[0],
              ).contiguous()
          output = torch.cat(
              tuple(tmp_tensor_list), comm_spec.gather_dim[1]
          ).contiguous()
  
      return output
  
imports: import operator
  from enum import Enum
  from functools import reduce
  import torch
  import torch.distributed
  from torch.distributed import ReduceOp
label_desc: Use image node_id 5 for calling a global function with example usage: _mix_gather(tensor, comm_spec) and returns: output