docstr: This method will find a path to transform source_spec to target_spec with
  a greedy algorithm.
  The basic idea is:
  Step1:
      Generate all one-step transform sequences from source_spec.
  Step2:
      Pick the 'best' sharding spec following the heuristic function.
  Step3:
      Repeat above steps until the source spec transform to target spec.
  
  During finding the transform path, communication cost will be accumulated, and it
  will be finally used in auto parallel solver.
  
  Additionally, to avoid repeating the path search in runtime, we cached all solved path
  in auto parallel strategy building time, which could handle most of cases in runtime.
  
  Argument:
      source_spec(ShardingSpec): ShardingSpec of the source activation.
      target_spec(ShardingSpec): ShardingSpec of the target activation.
  
  Return:
      transform_path(List[ShardingSpec]): The transform path from source_spec to target_spec,
                                          it contains the source_spec and target_spec.
      comm_action_sequence(List[CommSpec]): Keep the communication operations to complete the shape consistency in order.
      total_cost(float): total cost to complete shape consistency transform.
  
  Example:
      dim_partition_source = {1: [0, 1]}
      dim_partition_target = {0: [0, 1]}
      # DistSpec:
      #     shard_sequence: R,S01,R
      #     device_mesh_shape: (4, 4)
      sharding_spec_source = ShardingSpec(device_mesh, entire_shape, dim_partition_source)
      # DistSpec:
      #     shard_sequence: S01,R,R
      #     device_mesh_shape: (4, 4)
      sharding_spec_target = ShardingSpec(device_mesh, entire_shape, dim_partition_target)
      transform_path, comm_action_sequence, total_cost = shape_consistency_manager.shape_consistency(sharding_spec_source, sharding_spec_target)
      print(f'transform_path: {transform_path}')
      print(f'comm_action_sequence: {comm_action_sequence}')
      print(f'total_cost: {total_cost}')
  
  output:
      transform_path: [DistSpec:
              shard_sequence: R,S01,R
              device_mesh_shape: (4, 4), DistSpec:
              shard_sequence: R,S0,R
              device_mesh_shape: (4, 4), DistSpec:
              shard_sequence: S0,R,R
              device_mesh_shape: (4, 4), DistSpec:
              shard_sequence: S01,R,R
              device_mesh_shape: (4, 4)]
      comm_action_sequence: [CommSpec:(comm_pattern:allgather, gather_dim:1, logical_process_axis:1),
                             CommSpec:(comm_pattern:all2all, gather_dim:1, shard_dim:0, logical_process_axis: 0),
                             CommSpec:(comm_pattern:shard, shard_dim:0, logical_process_axis:1)]
      total_cost: 12294.402000000002
docstr_tok: ["This","method","will","find","a","path","to","transform","source_spec","to","target_spec","with","a","greedy","algorithm",".","The","basic","idea","is",":","Step1",":","Generate","all","one-step","transform","sequences","from","source_spec",".","Step2",":","Pick","the","'best","'","sharding","spec","following","the","heuristic","function",".","Step3",":","Repeat","above","steps","until","the","source","spec","transform","to","target","spec",".","During","finding","the","transform","path",",","communication","cost","will","be","accumulated",",","and","it","will","be","finally","used","in","auto","parallel","solver",".","Additionally",",","to","avoid","repeating","the","path","search","in","runtime",",","we","cached","all","solved","path","in","auto","parallel","strategy","building","time",",","which","could","handle","most","of","cases","in","runtime",".","Argument",":","source_spec","(","ShardingSpec",")",":","ShardingSpec","of","the","source","activation",".","target_spec","(","ShardingSpec",")",":","ShardingSpec","of","the","target","activation",".","Return",":","transform_path","(","List","[","ShardingSpec","]",")",":","The","transform","path","from","source_spec","to","target_spec",",","it","contains","the","source_spec","and","target_spec",".","comm_action_sequence","(","List","[","CommSpec","]",")",":","Keep","the","communication","operations","to","complete","the","shape","consistency","in","order",".","total_cost","(","float",")",":","total","cost","to","complete","shape","consistency","transform",".","Example",":","dim_partition_source","=","{","1",":","[","0",",","1","]","}","dim_partition_target","=","{","0",":","[","0",",","1","]","}","#","DistSpec",":","#","shard_sequence",":","R",",","S01",",","R","#","device_mesh_shape",":","(","4",",","4",")","sharding_spec_source","=","ShardingSpec","(","device_mesh",",","entire_shape",",","dim_partition_source",")","#","DistSpec",":","#","shard_sequence",":","S01",",","R",",","R","#","device_mesh_shape",":","(","4",",","4",")","sharding_spec_target","=","ShardingSpec","(","device_mesh",",","entire_shape",",","dim_partition_target",")","transform_path",",","comm_action_sequence",",","total_cost","=","shape_consistency_manager.shape_consistency","(","sharding_spec_source",",","sharding_spec_target",")","print","(","f'transform_path",":","{","transform_path","}","'",")","print","(","f'comm_action_sequence",":","{","comm_action_sequence","}","'",")","print","(","f'total_cost",":","{","total_cost","}","'",")","output",":","transform_path",":","[","DistSpec",":","shard_sequence",":","R",",","S01",",","R","device_mesh_shape",":","(","4",",","4",")",",","DistSpec",":","shard_sequence",":","R",",","S0",",","R","device_mesh_shape",":","(","4",",","4",")",",","DistSpec",":","shard_sequence",":","S0",",","R",",","R","device_mesh_shape",":","(","4",",","4",")",",","DistSpec",":","shard_sequence",":","S01",",","R",",","R","device_mesh_shape",":","(","4",",","4",")","]","comm_action_sequence",":","[","CommSpec",":","(","comm_pattern",":","allgather",",","gather_dim:1",",","logical_process_axis:1",")",",","CommSpec",":","(","comm_pattern",":","all2all",",","gather_dim:1",",","shard_dim:0",",","logical_process_axis",":","0",")",",","CommSpec",":","(","comm_pattern",":","shard",",","shard_dim:0",",","logical_process_axis:1",")","]","total_cost",":","12294.402000000002"]
code: def shape_consistency(
      self, source_spec: ShardingSpec, target_spec: ShardingSpec
  ) -> Tuple[List[ShardingSpec], List[CommSpec], float]:
      """
      This method will find a path to transform source_spec to target_spec with
      a greedy algorithm.
      The basic idea is:
      Step1:
          Generate all one-step transform sequences from source_spec.
      Step2:
          Pick the 'best' sharding spec following the heuristic function.
      Step3:
          Repeat above steps until the source spec transform to target spec.
  
      During finding the transform path, communication cost will be accumulated, and it
      will be finally used in auto parallel solver.
  
      Additionally, to avoid repeating the path search in runtime, we cached all solved path
      in auto parallel strategy building time, which could handle most of cases in runtime.
  
      Argument:
          source_spec(ShardingSpec): ShardingSpec of the source activation.
          target_spec(ShardingSpec): ShardingSpec of the target activation.
  
      Return:
          transform_path(List[ShardingSpec]): The transform path from source_spec to target_spec,
                                              it contains the source_spec and target_spec.
          comm_action_sequence(List[CommSpec]): Keep the communication operations to complete the shape consistency in order.
          total_cost(float): total cost to complete shape consistency transform.
  
      Example:
          dim_partition_source = {1: [0, 1]}
          dim_partition_target = {0: [0, 1]}
          # DistSpec:
          #     shard_sequence: R,S01,R
          #     device_mesh_shape: (4, 4)
          sharding_spec_source = ShardingSpec(device_mesh, entire_shape, dim_partition_source)
          # DistSpec:
          #     shard_sequence: S01,R,R
          #     device_mesh_shape: (4, 4)
          sharding_spec_target = ShardingSpec(device_mesh, entire_shape, dim_partition_target)
          transform_path, comm_action_sequence, total_cost = shape_consistency_manager.shape_consistency(sharding_spec_source, sharding_spec_target)
          print(f'transform_path: {transform_path}')
          print(f'comm_action_sequence: {comm_action_sequence}')
          print(f'total_cost: {total_cost}')
  
      output:
          transform_path: [DistSpec:
                  shard_sequence: R,S01,R
                  device_mesh_shape: (4, 4), DistSpec:
                  shard_sequence: R,S0,R
                  device_mesh_shape: (4, 4), DistSpec:
                  shard_sequence: S0,R,R
                  device_mesh_shape: (4, 4), DistSpec:
                  shard_sequence: S01,R,R
                  device_mesh_shape: (4, 4)]
          comm_action_sequence: [CommSpec:(comm_pattern:allgather, gather_dim:1, logical_process_axis:1),
                                 CommSpec:(comm_pattern:all2all, gather_dim:1, shard_dim:0, logical_process_axis: 0),
                                 CommSpec:(comm_pattern:shard, shard_dim:0, logical_process_axis:1)]
          total_cost: 12294.402000000002
      """
      MAX_TRANSFORM_STEPS = 20
      total_cost_dict = {"forward": 0, "backward": 0, "total": 0}
      total_steps = 0
      transform_path = []
      comm_action_sequence = []
      spec_pairs = (
          str(source_spec.sharding_sequence),
          str(target_spec.sharding_sequence),
      )
      self.cached_spec_pairs_transform_path[spec_pairs] = (None, None)
  
      # We do nothing if the sharding spec is all the same.
      if source_spec.sharding_sequence_difference(target_spec) == 0:
          self.cached_spec_pairs_transform_path[spec_pairs] = (
              transform_path,
              comm_action_sequence,
          )
          return (transform_path, comm_action_sequence, total_cost_dict)
  
      temp_sharding_spec = source_spec
  
      transform_path.append(temp_sharding_spec)
      # To avoid dead loop, the loop will break after MAX_TRANSFORM_STEPS transforms
      while total_steps <= MAX_TRANSFORM_STEPS:
          valid_transform_spec_dict = (
              self.get_all_one_step_transform_spec(
                  temp_sharding_spec, total_cost_dict
              )
          )
          best_difference_score = math.inf
  
          for (
              sharding_spec,
              info_pairs,
          ) in valid_transform_spec_dict.items():
              comm_spec, cost_dict = info_pairs
              spec_difference = (
                  sharding_spec.sharding_sequence_difference(
                      target_spec
                  )
              )
  
              if spec_difference == 0:
                  for phase, cost in total_cost_dict.items():
                      total_cost_dict[phase] = cost + cost_dict[phase]
                  transform_path.append(sharding_spec)
                  comm_action_sequence.append(comm_spec)
                  self.cached_spec_pairs_transform_path[spec_pairs] = (
                      transform_path,
                      comm_action_sequence,
                  )
                  return (
                      transform_path,
                      comm_action_sequence,
                      total_cost_dict,
                  )
  
              if spec_difference < best_difference_score:
                  temp_sharding_spec = sharding_spec
                  temp_cost_dict = cost_dict
                  temp_comm_spec = comm_spec
                  best_difference_score = spec_difference
  
          transform_path.append(temp_sharding_spec)
          comm_action_sequence.append(temp_comm_spec)
          for phase, cost in total_cost_dict.items():
              total_cost_dict[phase] = cost + temp_cost_dict[phase]
          total_steps += 1
  
      raise RuntimeError(
          f"Could not find a valid transform path with in {MAX_TRANSFORM_STEPS} steps."
      )
  
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 12 for calling the ShapeConsistencyManager obj's underlying member method code with example usage: obj.shape_consistency(source_spec, target_spec) and returns: transform_path,  comm_action_sequence,  total_cost_dict, transform_path,  comm_action_sequence,  total_cost_dict