code: def transpose_partition_dim(
      sharding_spec: ShardingSpec, dim1: int, dim2: int
  ) -> ShardingSpec:
      """
      Switch the sharding mesh dimensions for two tensor dimensions. This operation is in-place.
  
      Args:
          sharding_spec (ShardingSpec): the sharding spec for which partition dim are switched
          dim1 (int): the tensor dimension to switch
          dim2 (int): the tensor dimension to switch
      """
      assert (
          len(sharding_spec.entire_shape) >= 2
      ), "The entire_shape of the sharding spec must have at least 2 dimensions"
      dim_partition_dict = sharding_spec.dim_partition_dict
  
      # transpose the dim partition
      dim1_partition = dim_partition_dict.pop(dim1, None)
      dim2_partition = dim_partition_dict.pop(dim2, None)
  
      if dim1_partition:
          dim_partition_dict[dim2] = dim1_partition
      if dim2_partition:
          dim_partition_dict[dim1] = dim2_partition
  
      # get the transposed shape
      new_shape = list(sharding_spec.entire_shape[:])
      new_shape[dim2], new_shape[dim1] = (
          new_shape[dim1],
          new_shape[dim2],
      )
      new_shape = torch.Size(new_shape)
  
      # re-init the sharding spec
      sharding_spec.__init__(
          sharding_spec.device_mesh, new_shape, dim_partition_dict
      )
      return sharding_spec
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/utils/sharding.py
imports: import operator
  from copy import deepcopy
  from functools import reduce
  from typing import Dict
  import torch
  from colossalai.tensor.sharding_spec import ShardingSpec
label_desc: Use image node_id 1 for calling a global function with example usage: transpose_partition_dim(sharding_spec, dim1, dim2) and returns: sharding_spec