code: def infer_output_dim_partition_dict(
      input_dim_partition_dict: Dict[int, List[int]],
      reshape_mapping_dict: Dict[Tuple[int], Tuple[int]],
  ) -> Dict[Tuple[int], Tuple[int]]:
      """
      This method is used to infer the output dim partition dict for a reshape operation,
      given the input dim partition dict and reshape mapping dict.
      """
      assert check_keep_sharding_status(
          input_dim_partition_dict, reshape_mapping_dict
      ), "we only infer output dim partition dict for the reshape operation could keep sharding spec."
      sharded_dims = list(input_dim_partition_dict.keys())
      output_dim_partition_dict = {}
      for input_dims, output_dims in reshape_mapping_dict.items():
          for dim in input_dims:
              if dim in sharded_dims:
                  output_dim_partition_dict[
                      min(output_dims)
                  ] = input_dim_partition_dict[dim]
                  # we could break because input dims cannot contain two sharded dims, otherwise
                  # the keep sharding status check will fail.
                  break
      return output_dim_partition_dict
  
file_path: colossalai/colossalai/auto_parallel/tensor_shard/utils/reshape.py
imports: from enum import Enum
  from typing import Dict, List, Tuple
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: infer_output_dim_partition_dict(input_dim_partition_dict, reshape_mapping_dict) and returns: output_dim_partition_dict