code: def transpose(input: DTensor, dim0, dim1):
      shapes = transpose_shape(input.shardings.shapes, dim0, dim1)
      return DTensor(
          input.shardings.map_shard(
              lambda x: x.transpose(dim0, dim1).detach(),
              shapes=shapes.shapes,
              axis=shapes.axis,
          )
      )
  
label_desc: Use image node_id 1 for calling a global function with example usage: transpose(input, dim0, dim1) and returns: DTensor