[MLIR][TORCH] Add support for dynamic shape for Onnx.Transpose op (#2803)

Signed-Off By: Vivek Khandelwal <vivekkhandelwal1424@gmail.com>

Author: vivekkhandelwal1

lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp

@@ -1244,6 +1244,12 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
           current[i] = i;
         }
 
+        // Convert dynamic shape dimension.
+        for (unsigned i = 0; i < shape.size(); i++){
+          if (shape[i] == ShapedType::kDynamic)
+            shape[i] = Torch::kUnknownSize;
+        }
+
         for (int64_t i = 0; i < rank; ++i) {
           if (current[i] == permutations[i])
             continue;

test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir

@@ -968,6 +968,18 @@ func.func @test_transpose_all_permutations_4(%arg0: !torch.vtensor<[2,3,4],f32>)
   return %0 : !torch.vtensor<[4,2,3],f32>
 }
 
+// -----
+
+// CHECK-LABEL: func.func @test_transpose_dynamic
+func.func @test_transpose_dynamic(%arg0: !torch.vtensor<[?,32,5,128],f32>) -> !torch.vtensor<[?,5,32,128],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64} {
+  // CHECK-DAG:  %[[I1:.+]] = torch.constant.int 1
+  // CHECK-DAG:  %[[I2:.+]] = torch.constant.int 2
+  // CHECK:  %[[TRANSPOSE:.+]] = torch.aten.transpose.int %arg0, %[[I1]], %[[I2]] : !torch.vtensor<[?,32,5,128],f32>, !torch.int, !torch.int -> !torch.vtensor<[?,5,32,128],f32>
+  %0 = torch.operator "onnx.Transpose"(%arg0) {torch.onnx.perm = [0 : si64, 2 : si64, 1 : si64, 3 : si64]} : (!torch.vtensor<[?,32,5,128],f32>) -> !torch.vtensor<[?,5,32,128],f32>
+  return %0 : !torch.vtensor<[?,5,32,128],f32>
+}
+
+
 // -----
 
 // CHECK-LABEL: func.func @test_slice