class_docstr: Communication spec is used to record the communication action. It has two main functions:
  1. Compute the communication cost which will be used in auto parallel solver.
  2. Convert the communication spec to real action which will be used in runtime.
  It contains comm_pattern to determine the
  communication method, sharding_spec to determine the communication size, gather_dim and shard_dim
  to determine the buffer shape, and logical_process_axis
  
  Argument:
      comm_pattern(CollectiveCommPattern): describe the communication method used in this spec.
      sharding_spec(ShardingSpec): This is sharding spec of the tensor which will join the communication action.
      gather_dim(int, Optional): The gather_dim of the tensor will be gathered.
      shard_dim(int, Optional): The shard_dim of the tensor will be sharded.
      logical_process_axis(Union(int, List[int]), Optional): The mesh_dim to implement the communication action.
class_docstr_tok: ["Communication","spec","is","used","to","record","the","communication","action",".","It","has","two","main","functions",":","1",".","Compute","the","communication","cost","which","will","be","used","in","auto","parallel","solver",".","2",".","Convert","the","communication","spec","to","real","action","which","will","be","used","in","runtime",".","It","contains","comm_pattern","to","determine","the","communication","method",",","sharding_spec","to","determine","the","communication","size",",","gather_dim","and","shard_dim","to","determine","the","buffer","shape",",","and","logical_process_axis","Argument",":","comm_pattern","(","CollectiveCommPattern",")",":","describe","the","communication","method","used","in","this","spec",".","sharding_spec","(","ShardingSpec",")",":","This","is","sharding","spec","of","the","tensor","which","will","join","the","communication","action",".","gather_dim","(","int",",","Optional",")",":","The","gather_dim","of","the","tensor","will","be","gathered",".","shard_dim","(","int",",","Optional",")",":","The","shard_dim","of","the","tensor","will","be","sharded",".","logical_process_axis","(","Union","(","int",",","List","[","int","]",")",",","Optional",")",":","The","mesh_dim","to","implement","the","communication","action","."]
code: def __init__(
      self,
      comm_pattern,
      sharding_spec,
      gather_dim=None,
      shard_dim=None,
      logical_process_axis=None,
      forward_only=False,
      mix_gather=False,
  ):
      self.comm_pattern = comm_pattern
      self.sharding_spec = sharding_spec
      self.gather_dim = gather_dim
      self.shard_dim = shard_dim
      self.logical_process_axis = logical_process_axis
      self.forward_only = forward_only
      if isinstance(self.logical_process_axis, list):
          if not mix_gather:
              self.device_mesh = (
                  self.sharding_spec.device_mesh.flatten()
              )
              self.logical_process_axis = 0
          else:
              self.device_meshes = (
                  self.sharding_spec.device_mesh.flatten_device_meshes
              )
              self.device_mesh = (
                  self.sharding_spec.device_mesh.flatten_device_mesh
              )
              # Create a new member `logical_process_axes` to distinguish from original flatten
              self.logical_process_axes = logical_process_axis
      else:
          self.device_mesh = self.sharding_spec.device_mesh
  
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 1 to create a new CommSpec object  with example: obj = CommSpec(comm_pattern, sharding_spec, gather_dim, shard_dim, logical_process_axis, forward_only, mix_gather)