code: def transpose_shape(shape: _ShardingShapes, dim0, dim1):
      # transpose shapes
      shapes = []
      for s in shape.shapes:
          s = list(s)
          s[dim0], s[dim1] = s[dim1], s[dim0]
          shapes.append(torch.Size(s))
      # transpose axis
      axis = shape.axis
      if dim0 == axis:
          axis = dim1
      elif dim1 == axis:
          axis = dim0
      return _ShardingShapes(shapes=shapes, axis=axis)
  
label_desc: Use image node_id 2 for calling a global function with example usage: transpose_shape(shape, dim0, dim1) and returns: _ShardingShapes