docstr: Get all valid sharding specs from source_spec with single all-to-all operation, and
  accumulate communication cost on origin cost which will finally be used in auto sharding solver.
  For the all-to-all operation, we just care about the pairs containing 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-to-all 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_to_all_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: R,S1,S0
      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","all-to-all","operation",",","and","accumulate","communication","cost","on","origin","cost","which","will","finally","be","used","in","auto","sharding","solver",".","For","the","all-to-all","operation",",","we","just","care","about","the","pairs","containing","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-to-all","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_to_all_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",":","R",",","S1",",","S0","device_mesh_shape",":","(","4",",","4",")",":","0",",","DistSpec",":","shard_sequence",":","S0",",","R",",","S1","device_mesh_shape",":","(","4",",","4",")",":","0","}"]
code: def get_all_all_to_all_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-to-all operation, and
      accumulate communication cost on origin cost which will finally be used in auto sharding solver.
      For the all-to-all operation, we just care about the pairs containing 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-to-all 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_to_all_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: R,S1,S0
          device_mesh_shape: (4, 4): 0, DistSpec:
          shard_sequence: S0,R,S1
          device_mesh_shape: (4, 4): 0}
      """
      valid_spec_dict = {}
      comm_pattern = CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD
      tensor_dims = len(source_spec.entire_shape)
      for f_index in range(tensor_dims - 1):
          for b_index in range(f_index + 1, tensor_dims):
              # skip (R, R) cases
              if (
                  f_index not in source_spec.dim_partition_dict
                  and b_index not in source_spec.dim_partition_dict
              ):
                  continue
              else:
                  if f_index in source_spec.dim_partition_dict:
                      # skip (S01, R) -> (R, S01) is NOT allowed
                      if (
                          len(source_spec.dim_partition_dict[f_index])
                          >= 2
                      ):
                          continue
                      f_target_pair = (
                          f_index,
                          deepcopy(
                              source_spec.dim_partition_dict[f_index]
                          ),
                      )
                  else:
                      f_target_pair = (f_index, [])
                  if b_index in source_spec.dim_partition_dict:
                      # skip (R, S01) -> (S01, R) is NOT allowed
                      if (
                          len(source_spec.dim_partition_dict[b_index])
                          >= 2
                      ):
                          continue
                      b_target_pair = (
                          b_index,
                          deepcopy(
                              source_spec.dim_partition_dict[b_index]
                          ),
                      )
                  else:
                      b_target_pair = (b_index, [])
  
              # skip (S1, S0) -> S10
              if (
                  f_target_pair[1]
                  and b_target_pair[1]
                  and f_target_pair[1][0] >= b_target_pair[1][0]
              ):
                  continue
              f_shard_list, b_shard_list = all_to_all_simulator(
                  f_target_pair, b_target_pair
              )
              f_index = f_target_pair[0]
              b_index = b_target_pair[0]
  
              # generate the CommSpec to record the action of source_sharding_spec->new_sharding_spec
              if len(f_shard_list) < len(f_target_pair[1]):
                  gather_dim = f_index
                  shard_dim = b_index
                  logical_process_axis = f_target_pair[1][-1]
              else:
                  gather_dim = b_index
                  shard_dim = f_index
                  logical_process_axis = b_target_pair[1][-1]
              comm_spec = CommSpec(
                  comm_pattern,
                  sharding_spec=source_spec,
                  gather_dim=gather_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()
              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 f_shard_list:
                  new_dim_partition_dict[f_index] = f_shard_list
              else:
                  new_dim_partition_dict.pop(f_index)
              if b_shard_list:
                  new_dim_partition_dict[b_index] = b_shard_list
              else:
                  new_dim_partition_dict.pop(b_index)
  
              # 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 7 for calling the ShapeConsistencyManager obj's underlying member method code with example usage: obj.get_all_all_to_all_spec(source_spec, orig_cost_dict) and returns: valid_spec_dict