class_docstr: A logical view of a physical cluster. For example, we could view a physical cluster
  with 16 devices as a device mesh with shape (2, 2, 4) or (4, 4).
  
  Arguments:
      physical_mesh_id (torch.Tensor): physical view of the devices in global rank.
      logical_mesh_id (torch.Tensor): logical view of the devices in global rank.
      mesh_shape (torch.Size, optional): shape of logical view.
      mesh_alpha (List[float], optional): coefficients used for computing
          communication cost (default: None)
      mesh_beta (List[float], optional): coefficients used for computing
          communication cost (default: None)
      init_process_group (bool, optional): initialize logical process group
          during initializing the DeviceMesh instance if the init_process_group set to True.
          Otherwise, users need to call create_process_groups_for_logical_mesh manually to init logical process group.
          (default: False)
      device (str): the device for the process groups used by the DeviceMesh instance. (default: 'cuda')
docstr: Give a global rank and return all global ranks involved in its associated process group in each axis.
  
  Example:
  
  ```python
  physical_mesh_id = torch.arange(0, 16)
  mesh_shape = (4, 4)
  
  # logical mesh will look like
  # [[0, 1, 2, 3],
  #  [4, 5, 6, 7],
  #  [8, 9, 10,11],
  #  [12,13,14,15]]
  
  device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
  print(device_mesh.collate_global_ranks_in_same_process_group(0))
  
  # key is axis name
  # value is a list of global ranks in same axis with rank 0
  # output will look like
  # {
      0: [0, 4, 8, 12],
      1: [0, 1, 2, 3]
  #  }
docstr_tok: ["Give","a","global","rank","and","return","all","global","ranks","involved","in","its","associated","process","group","in","each","axis",".","Example",":","``","`","python","physical_mesh_id","=","torch.arange","(","0",",","16",")","mesh_shape","=","(","4",",","4",")","#","logical","mesh","will","look","like","#","[","[","0",",","1",",","2",",","3","]",",","#","[","4",",","5",",","6",",","7","]",",","#","[","8",",","9",",","10,11","]",",","#","[","12,13,14,15","]","]","device_mesh","=","DeviceMesh","(","physical_mesh_id",",","mesh_shape",")","print","(","device_mesh.collate_global_ranks_in_same_process_group","(","0",")",")","#","key","is","axis","name","#","value","is","a","list","of","global","ranks","in","same","axis","with","rank","0","#","output","will","look","like","#","{","0",":","[","0",",","4",",","8",",","12","]",",","1",":","[","0",",","1",",","2",",","3","]","#","}"]
code: def _collate_global_ranks_in_same_process_group(self, global_rank):
      """
      Give a global rank and return all global ranks involved in its associated process group in each axis.
  
      Example:
  
      ```python
      physical_mesh_id = torch.arange(0, 16)
      mesh_shape = (4, 4)
  
      # logical mesh will look like
      # [[0, 1, 2, 3],
      #  [4, 5, 6, 7],
      #  [8, 9, 10,11],
      #  [12,13,14,15]]
  
      device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
      print(device_mesh.collate_global_ranks_in_same_process_group(0))
  
      # key is axis name
      # value is a list of global ranks in same axis with rank 0
      # output will look like
      # {
          0: [0, 4, 8, 12],
          1: [0, 1, 2, 3]
      #  }
      """
      # We have init the global rank to local rank by calling _init_global_to_logical_rank_mapping
      # for self._global_to_local_rank_mapping
      # the key is the global rank
      # the value is the list of local ranks corresponding to the global rank with respect of different axes
      # we can see the list of local ranks as the process coordinates for simplicity
      # the key and value are all unique, therefore,
      # we can also to use the coordinates to find the global rank
  
      # =========================================================================
      # Step 1
      # find all the process_coordinates for processes in the same process group
      # as the given global rank
      # =========================================================================
  
      # each
      processes_in_the_same_process_group = {}
  
      for dim in range(self.logical_mesh_id.dim()):
          # iterate over the dimension size so that we can include all processes
          # in the same process group in the given axis
          # the _local_rank refers to the local rank of the current process
          for _local_rank in range(self.logical_mesh_id.shape[dim]):
              # if this dimension is not initialized yet,
              # initialize it with an empty array
              if dim not in processes_in_the_same_process_group:
                  processes_in_the_same_process_group[dim] = []
  
              # get the local rank corresponding to the global rank
              process_coordinates = self._global_to_local_rank_mapping[
                  global_rank
              ].copy()
  
              # replace the local rank in the given dimension with the
              # local rank of the current process iterated
              process_coordinates[dim] = _local_rank
              processes_in_the_same_process_group[dim].append(
                  process_coordinates
              )
  
      # =================================================================
      # Step 2
      # Use local rank combination to find its corresponding global rank
      # =================================================================
      # the key of the dict is the axis
      # the value is the list of global ranks which are in the same process group as the given global rank
      global_pg_ranks = {}
      for (
          dim,
          coordinates_of_all_processes,
      ) in processes_in_the_same_process_group.items():
          global_pg_ranks[dim] = []
          for process_coordinates in coordinates_of_all_processes:
              # find the global rank by local rank combination
              for (
                  _global_rank,
                  _process_coordinates,
              ) in self._global_to_local_rank_mapping.items():
                  if process_coordinates == _process_coordinates:
                      global_pg_ranks[dim].append(_global_rank)
      return global_pg_ranks
  
imports: import operator
  from dataclasses import dataclass
  from functools import reduce
  from typing import Dict, List, Union
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
label_desc: Use image node_id 15 for calling the DeviceMesh obj's underlying member method code with example usage: obj._collate_global_ranks_in_same_process_group(global_rank) and returns: global_pg_ranks