Minor fix for negative axes

Add negative axes lit tests for MVN

Signed-off-by: Zahid Wakeel <[email protected]>

Author: zahidwx

lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

@@ -1641,7 +1641,7 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
           if (!Torch::isValidDim(dim, inputRank)) {
             return failure();
           }
-          reduced_shape[i] = 1;
+          reduced_shape[dim] = 1;
         }
         Torch::ValueTensorType reducedOutTy = Torch::ValueTensorType::get(
             resultType.getContext(), reduced_shape, resultType.getDtype());

test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir

@@ -1646,6 +1646,30 @@ func.func @test_meanvarnorm_axes(%arg0: !torch.vtensor<[3,5,2,2],f32>) -> !torch
 
 // -----
 
+// CHECK-LABEL:   func.func @test_meanvarnorm_neg_axes(
+func.func @test_meanvarnorm_neg_axes(%arg0: !torch.vtensor<[3,5,2,2],f32>) -> !torch.vtensor<[3,5,2,2],f32> attributes {torch.onnx_meta.ir_version = 3 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
+    // CHECK-SAME:      %[[ARG0:.*]]: !torch.vtensor<[3,5,2,2],f32>) -> !torch.vtensor<[3,5,2,2],f32> attributes {torch.onnx_meta.ir_version = 3 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
+    // CHECK:           %[[VAL_0:.*]] = torch.constant.bool true
+    // CHECK:           %[[VAL_1:.*]] = torch.constant.bool false
+    // CHECK:           %[[VAL_2:.*]] = torch.constant.none
+    // CHECK:           %[[VAL_3:.*]] = torch.constant.int -1
+    // CHECK:           %[[VAL_4:.*]] = torch.constant.int -3
+    // CHECK:           %[[VAL_5:.*]] = torch.prim.ListConstruct %[[VAL_3]], %[[VAL_4]] : (!torch.int, !torch.int) -> !torch.list<int>
+    // CHECK:           %[[VAL_6:.*]] = torch.aten.mean.dim %[[ARG0]], %[[VAL_5]], %[[VAL_0]], %[[VAL_2]] : !torch.vtensor<[3,5,2,2],f32>, !torch.list<int>, !torch.bool, !torch.none -> !torch.vtensor<[3,1,2,1],f32>
+    // CHECK:           %[[VAL_7:.*]] = torch.aten.var.dim %[[ARG0]], %[[VAL_5]], %[[VAL_1]], %[[VAL_0]] : !torch.vtensor<[3,5,2,2],f32>, !torch.list<int>, !torch.bool, !torch.bool -> !torch.vtensor<[3,1,2,1],f32>
+    // CHECK:           %[[VAL_8:.*]] = torch.constant.int 1
+    // CHECK:           %[[VAL_9:.*]] = torch.constant.float 1.000000e-09
+    // CHECK:           %[[VAL_10:.*]] = torch.aten.add.Scalar %[[VAL_7]], %[[VAL_9]], %[[VAL_8]] : !torch.vtensor<[3,1,2,1],f32>, !torch.float, !torch.int -> !torch.vtensor<[3,1,2,1],f32>
+    // CHECK:           %[[VAL_11:.*]] = torch.aten.sqrt %[[VAL_10]] : !torch.vtensor<[3,1,2,1],f32> -> !torch.vtensor<[3,1,2,1],f32>
+    // CHECK:           %[[VAL_12:.*]] = torch.aten.sub.Tensor %[[ARG0]], %[[VAL_6]], %[[VAL_8]] : !torch.vtensor<[3,5,2,2],f32>, !torch.vtensor<[3,1,2,1],f32>, !torch.int -> !torch.vtensor<[3,5,2,2],f32>
+    // CHECK:           %[[VAL_13:.*]] = torch.aten.div.Tensor %[[VAL_12]], %[[VAL_11]] : !torch.vtensor<[3,5,2,2],f32>, !torch.vtensor<[3,1,2,1],f32> -> !torch.vtensor<[3,5,2,2],f32>
+    // CHECK:           return %[[VAL_13]] : !torch.vtensor<[3,5,2,2],f32>
+    // CHECK:         }
+  %0 = torch.operator "onnx.MeanVarianceNormalization"(%arg0) {torch.onnx.axes = [-1 : si64, -3 : si64]} : (!torch.vtensor<[3,5,2,2],f32>) -> !torch.vtensor<[3,5,2,2],f32>
+  return %0 : !torch.vtensor<[3,5,2,2],f32>
+
+// -----
+
 // CHECK-LABEL: func.func @test_not_2d
 func.func @test_not_2d(%arg0: !torch.vtensor<[3,4],i1>) -> !torch.vtensor<[3,4],i1> attributes {torch.onnx_meta.ir_version = 3 : si64, torch.onnx_meta.opset_version = 1 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
     // CHECK: torch.aten.bitwise_not %arg0 : !torch.vtensor<[3,4],i1> -> !torch.vtensor<[3,4],i1>