docstr: Get all valid sharding specs from source_spec with single shard operation, and
  accumulate communication cost on origin cost which will finally be used in auto sharding solver.
  For the sharding operation, we just care about legal sharding dimensions.
  
  Argument:
      source_spec(ShardingSpec): the ShardingSpec of the source_spec.
      orig_cost(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-to-all operation.
  
  Example:
      dim_partition_dict = {0: [0]}
      # DistSpec:
      #     shard_sequence: S0,R,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_shard_spec(sharding_spec, {'forward': 0, 'backward': 0, 'total': 0})
      print(rst_dict)
  
  Output:
      {DistSpec:
      shard_sequence: S01,R,R
      device_mesh_shape: (4, 4): 0, DistSpec:
      shard_sequence: S0,S1,R
      device_mesh_shape: (4, 4): 0, DistSpec:
      shard_sequence: S0,R,S1
      device_mesh_shape: (4, 4): 0}
docstr_tok: ["Get","all","valid","sharding","specs","from","source_spec","with","single","shard","operation",",","and","accumulate","communication","cost","on","origin","cost","which","will","finally","be","used","in","auto","sharding","solver",".","For","the","sharding","operation",",","we","just","care","about","legal","sharding","dimensions",".","Argument",":","source_spec","(","ShardingSpec",")",":","the","ShardingSpec","of","the","source_spec",".","orig_cost","(","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-to-all","operation",".","Example",":","dim_partition_dict","=","{","0",":","[","0","]","}","#","DistSpec",":","#","shard_sequence",":","S0",",","R",",","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_shard_spec","(","sharding_spec",",","{","'forward","'",":","0",",","'backward","'",":","0",",","'total","'",":","0","}",")","print","(","rst_dict",")","Output",":","{","DistSpec",":","shard_sequence",":","S01",",","R",",","R","device_mesh_shape",":","(","4",",","4",")",":","0",",","DistSpec",":","shard_sequence",":","S0",",","S1",",","R","device_mesh_shape",":","(","4",",","4",")",":","0",",","DistSpec",":","shard_sequence",":","S0",",","R",",","S1","device_mesh_shape",":","(","4",",","4",")",":","0","}"]
code: def get_all_shard_spec(
      self, source_spec: ShardingSpec, orig_cost_dict
  ):
      """
      Get all valid sharding specs from source_spec with single shard operation, and
      accumulate communication cost on origin cost which will finally be used in auto sharding solver.
      For the sharding operation, we just care about legal sharding dimensions.
  
      Argument:
          source_spec(ShardingSpec): the ShardingSpec of the source_spec.
          orig_cost(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-to-all operation.
  
      Example:
          dim_partition_dict = {0: [0]}
          # DistSpec:
          #     shard_sequence: S0,R,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_shard_spec(sharding_spec, {'forward': 0, 'backward': 0, 'total': 0})
          print(rst_dict)
  
      Output:
          {DistSpec:
          shard_sequence: S01,R,R
          device_mesh_shape: (4, 4): 0, DistSpec:
          shard_sequence: S0,S1,R
          device_mesh_shape: (4, 4): 0, DistSpec:
          shard_sequence: S0,R,S1
          device_mesh_shape: (4, 4): 0}
      """
      valid_spec_dict = {}
      comm_pattern = CollectiveCommPattern.SPLIT_FWD_GATHER_BWD
  
      # legal sharding dims means the mesh_id is still available to use.
      legal_sharding_dims = [
          i for i in range(len(source_spec.device_mesh.shape))
      ]
      for dim, shard_list in source_spec.dim_partition_dict.items():
          for element in shard_list:
              legal_sharding_dims.remove(element)
      if len(legal_sharding_dims) == 0:
          return valid_spec_dict
  
      tensor_dims = len(source_spec.entire_shape)
  
      for index in range(tensor_dims):
          if index not in source_spec.dim_partition_dict:
              shard_list_list = shard_simulator(
                  (index, []), legal_sharding_dims
              )
          else:
              shard_list_list = shard_simulator(
                  (index, source_spec.dim_partition_dict[index]),
                  legal_sharding_dims,
              )
          if not shard_list_list:
              continue
          for shard_list in shard_list_list:
              new_dim_partition_dict = deepcopy(
                  source_spec.dim_partition_dict
              )
              new_dim_partition_dict[index] = shard_list
  
              # generate the CommSpec to record the action of source_sharding_spec->new_sharding_spec
              shard_dim = index
              logical_process_axis = shard_list[-1]
              comm_spec = CommSpec(
                  comm_pattern,
                  sharding_spec=source_spec,
                  gather_dim=shard_dim,
                  shard_dim=shard_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 8 for calling the ShapeConsistencyManager obj's underlying member method code with example usage: obj.get_all_shard_spec(source_spec, orig_cost_dict) and returns: valid_spec_dict, valid_spec_dict