code: def test_shape_consistency():
      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
      )
  
      transform_path_str = "->".join(
          [
              str(sharding_spec.sharding_sequence)
              for sharding_spec in transform_path
          ]
      )
      assert (
          transform_path_str
          == "[R, S01, R]->[R, S0, R]->[S0, R, R]->[S01, R, R]"
      )
  
      # all-gather(S01) -> S0
      assert (
          comm_action_sequence[0].comm_pattern
          == CollectiveCommPattern.GATHER_FWD_SPLIT_BWD
      )
      assert comm_action_sequence[0].gather_dim == 1
      assert comm_action_sequence[0].logical_process_axis == 1
  
      # all-to-all(R, S0) -> [S0, R]
      assert (
          comm_action_sequence[1].comm_pattern
          == CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD
      )
      assert comm_action_sequence[1].gather_dim == 1
      assert comm_action_sequence[1].shard_dim == 0
      assert comm_action_sequence[1].logical_process_axis == 0
  
      # shard(S0) -> [S01]
      assert (
          comm_action_sequence[2].comm_pattern
          == CollectiveCommPattern.SPLIT_FWD_GATHER_BWD
      )
      assert comm_action_sequence[2].shard_dim == 0
      assert comm_action_sequence[2].logical_process_axis == 1
  
      assert (
          shape_consistency_manager.cached_spec_pairs_transform_path[
              ("[R, S01, R]", "[S01, R, R]")
          ][0]
          == transform_path
      )
      assert (
          shape_consistency_manager.cached_spec_pairs_transform_path[
              ("[R, S01, R]", "[S01, R, R]")
          ][1]
          == comm_action_sequence
      )
  
imports: import torch
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.tensor.shape_consistency import CollectiveCommPattern, ShapeConsistencyManager
  from colossalai.tensor.sharding_spec import ShardingSpec
label_desc: Use image node_id 2 for calling a global function with example usage: test_shape_consistency() without return types