docstr: Get all valid sharding specs from source_spec with single all-gather operation, and
  accumulate communication cost on origin cost which will finally be used in auto sharding solver.
  For the all-gather operation, we just care about the S dimension.
  
  Argument:
      source_spec(ShardingSpec): the ShardingSpec of the source_spec.
      orig_cost(Dict[str, float]): the original communication cost before this operation.
  
  Return:
      valid_spec_dict(Dict[ShardingSpec, float]): all valid sharding specs from source_spec with single all-gather operation.
  
  Example:
      dim_partition_dict = {0: [0], 1: [1]}
      # DistSpec:
      #     shard_sequence: S0,S1,R
      #     device_mesh_shape: (4, 4)
      sharding_spec = ShardingSpec(device_mesh, entire_shape, dim_partition_dict)
      shape_consistency_manager = ShapeConsistencyManager()
      rst_dict = shape_consistency_manager.get_all_all_gather_spec(sharding_spec, {'forward': 0, 'backward': 0, 'total': 0})
      print(rst_dict)
  
  Output:
      {DistSpec:
      shard_sequence: R,S1,R
      device_mesh_shape: (4, 4): 0, DistSpec:
      shard_sequence: S0,R,R
      device_mesh_shape: (4, 4): 0}
docstr_tok: ["Get","all","valid","sharding","specs","from","source_spec","with","single","all-gather","operation",",","and","accumulate","communication","cost","on","origin","cost","which","will","finally","be","used","in","auto","sharding","solver",".","For","the","all-gather","operation",",","we","just","care","about","the","S","dimension",".","Argument",":","source_spec","(","ShardingSpec",")",":","the","ShardingSpec","of","the","source_spec",".","orig_cost","(","Dict","[","str",",","float","]",")",":","the","original","communication","cost","before","this","operation",".","Return",":","valid_spec_dict","(","Dict","[","ShardingSpec",",","float","]",")",":","all","valid","sharding","specs","from","source_spec","with","single","all-gather","operation",".","Example",":","dim_partition_dict","=","{","0",":","[","0","]",",","1",":","[","1","]","}","#","DistSpec",":","#","shard_sequence",":","S0",",","S1",",","R","#","device_mesh_shape",":","(","4",",","4",")","sharding_spec","=","ShardingSpec","(","device_mesh",",","entire_shape",",","dim_partition_dict",")","shape_consistency_manager","=","ShapeConsistencyManager","(",")","rst_dict","=","shape_consistency_manager.get_all_all_gather_spec","(","sharding_spec",",","{","'forward","'",":","0",",","'backward","'",":","0",",","'total","'",":","0","}",")","print","(","rst_dict",")","Output",":","{","DistSpec",":","shard_sequence",":","R",",","S1",",","R","device_mesh_shape",":","(","4",",","4",")",":","0",",","DistSpec",":","shard_sequence",":","S0",",","R",",","R","device_mesh_shape",":","(","4",",","4",")",":","0","}"]
code: def get_all_all_gather_spec(
      self, source_spec: ShardingSpec, orig_cost_dict: Dict[str, float]
  ) -> Dict[ShardingSpec, float]:
      """
      Get all valid sharding specs from source_spec with single all-gather operation, and
      accumulate communication cost on origin cost which will finally be used in auto sharding solver.
      For the all-gather operation, we just care about the S dimension.
  
      Argument:
          source_spec(ShardingSpec): the ShardingSpec of the source_spec.
          orig_cost(Dict[str, float]): the original communication cost before this operation.
  
      Return:
          valid_spec_dict(Dict[ShardingSpec, float]): all valid sharding specs from source_spec with single all-gather operation.
  
      Example:
          dim_partition_dict = {0: [0], 1: [1]}
          # DistSpec:
          #     shard_sequence: S0,S1,R
          #     device_mesh_shape: (4, 4)
          sharding_spec = ShardingSpec(device_mesh, entire_shape, dim_partition_dict)
          shape_consistency_manager = ShapeConsistencyManager()
          rst_dict = shape_consistency_manager.get_all_all_gather_spec(sharding_spec, {'forward': 0, 'backward': 0, 'total': 0})
          print(rst_dict)
  
      Output:
          {DistSpec:
          shard_sequence: R,S1,R
          device_mesh_shape: (4, 4): 0, DistSpec:
          shard_sequence: S0,R,R
          device_mesh_shape: (4, 4): 0}
      """
      valid_spec_dict = {}
      comm_pattern = CollectiveCommPattern.GATHER_FWD_SPLIT_BWD
      for target_pair in source_spec.dim_partition_dict.items():
          shard_list = all_gather_simulator(target_pair)
          index = target_pair[0]
          new_dim_partition_dict = deepcopy(
              source_spec.dim_partition_dict
          )
  
          # We won't add empty list into dim_partition_dict
          # The key will be popped if the related shard_list is empty
          if shard_list:
              new_dim_partition_dict[index] = shard_list
          else:
              new_dim_partition_dict.pop(index)
  
          # generate the CommSpec to record the action of source_sharding_spec->new_sharding_spec
          gather_dim = index
          logical_process_axis = target_pair[1][-1]
          comm_spec = CommSpec(
              comm_pattern,
              sharding_spec=source_spec,
              gather_dim=gather_dim,
              # shard_dim will be used during backward
              shard_dim=gather_dim,
              logical_process_axis=logical_process_axis,
              forward_only=self.forward_only,
          )
  
          # compute the communication cost with CommSpec
          cost_dict = comm_spec.get_comm_cost()
  
          # generate new sharding spec
          try:
              new_sharding_spec = ShardingSpec(
                  source_spec.device_mesh,
                  source_spec.entire_shape,
                  dim_partition_dict=new_dim_partition_dict,
              )
              for phase, cost in cost_dict.items():
                  cost_dict[phase] = cost + orig_cost_dict[phase]
              valid_spec_dict[new_sharding_spec] = (
                  comm_spec,
                  cost_dict,
              )
          except ShardingSpecException:
              pass
      return valid_spec_dict
  
imports: import math
  from copy import deepcopy
  from dataclasses import dataclass
  from typing import Dict, List, Tuple
  import numpy
  import torch
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, TrainCycleItem
  from colossalai.context.singleton_meta import SingletonMeta
  from colossalai.tensor.sharding_spec import ShardingSpec, ShardingSpecException
  from colossalai.tensor.utils import all_gather_simulator, all_to_all_simulator, mix_gather_simulator, shard_simulator
  from .comm_spec import *
label_desc: Use image node_id 6 for calling the ShapeConsistencyManager obj's underlying member method code with example usage: obj.get_all_all_gather_spec(source_spec, orig_cost_dict) and returns: valid_spec_dict