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')
class_docstr_tok: ["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","'",")"]
code: def __init__(
      self,
      physical_mesh_id: torch.Tensor,
      mesh_shape: torch.Size = None,
      logical_mesh_id: torch.Tensor = None,
      mesh_alpha: List[float] = None,
      mesh_beta: List[float] = None,
      init_process_group: bool = False,
      device: str = "cuda",
  ):
      # ============================
      # Physical & Logical Mesh IDs
      # ============================
      self._physical_mesh_id = physical_mesh_id
      assert (
          physical_mesh_id.dim() == 1
      ), "physical_mesh_id should be a 1D tensor."
  
      # logical mesh ids can be obtained via two ways
      # 1. provide physical mesh id and provide mesh shape
      # 2. directly supply the logical mesh id
      assert mesh_shape is None or logical_mesh_id is None, (
          "Only one of mesh_shape and logical_mesh_id can be specified."
          "Logical mesh IDs are obtained from either mesh_shape + physical_mesh_id or directly from the user-supplied logical_mesh_id"
      )
  
      if logical_mesh_id is None:
          self._mesh_shape = mesh_shape
          self._logical_mesh_id = self._physical_mesh_id.reshape(
              self._mesh_shape
          )
      else:
          self._logical_mesh_id = logical_mesh_id
          self._mesh_shape = self._logical_mesh_id.shape
  
      # ensure two things:
      # 1. logical and physical mesh IDs should contain the same elements
      # 2. there is no duplicate IDs in each mesh, e.g. [2, 2] is not allowed
      assert torch.equal(
          torch.unique(self._physical_mesh_id),
          torch.unique(self.logical_mesh_id),
      ), "physical and logical mesh IDs should contain the same elements, please check if you have consistent physical_mesh_id and logical_mesh_id."
      assert (
          torch.unique(self._physical_mesh_id).numel()
          == self._physical_mesh_id.numel()
      ), "Found duplicate IDs in the physical_mesh_id and this is not allowed, please check your physical_mesh_id again."
      assert (
          torch.unique(self.logical_mesh_id).numel()
          == self.logical_mesh_id.numel()
      ), "Found duplicate IDs in the logical_mesh_id and this is not allowed, please check your logical_mesh_id again."
  
      # ===============================================
      # coefficient for alpha-beta communication model
      # alpha is latency and beta is bandwidth
      # ===============================================
      # if the values are not provided, we assume they are 1 for simplicity
      if mesh_alpha is None:
          mesh_alpha = [1] * len(self._mesh_shape)
      if mesh_beta is None:
          mesh_beta = [1] * len(self._mesh_shape)
  
      self.mesh_alpha = tuple(mesh_alpha)
      self.mesh_beta = tuple(mesh_beta)
  
      # ensure the alpha and beta have the same shape
      assert len(self.mesh_alpha) == len(
          self.mesh_beta
      ), "mesh_alpha and mesh_beta should have the same length, please check your mesh_alpha and mesh_beta again."
  
      # =========================
      # Device for Process Group
      # =========================
      self._device = device
      self._dist_backend = self._DIST_BACKEND[device]
  
      # =========================
      # Process Group Management
      # =========================
      # the _global_to_local_rank_mapping is structured as follows
      # {
      #    <global-rank>: [ <local-rank-on-axis-0>, <local-rank-on-axis-1>, <local-rank-on-axis-2>, ...]
      # }
      self._global_to_local_rank_mapping = dict()
      self._init_global_to_logical_rank_mapping(
          mapping=self._global_to_local_rank_mapping,
          tensor=self.logical_mesh_id,
      )
  
      # create process group
      self._process_group_dict = {}
      self._ranks_in_the_process_group = {}
      self._global_rank_of_current_process = None
      self._is_initialized = False
  
      # attribute used to indicate whether this object
      # is created using DeviceMesh.from_process_group
      # this attribute can be used to do some check in methods
      # such get_process_group as no global rank information
      # is known if created with from_process_group
      self._is_init_from_process_group = False
  
      # initialize process group if specified
      self._init_ranks_in_the_same_group()
      self._init_process_group = init_process_group
      if init_process_group:
          self.init_logical_process_group()
  
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 1 to create a new DeviceMesh object  with example: obj = DeviceMesh(physical_mesh_id, mesh_shape, logical_mesh_id, mesh_alpha, mesh_beta, init_process_group, device)