Lower to torch dialect without expansion

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

Author: zahidwx

lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

@@ -1602,6 +1602,67 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
             /* cudnn enabled */ boolFalse);
         return success();
       });
+  patterns.onOp(
+      "MeanVarianceNormalization", 13,
+      [](OpBinder binder, ConversionPatternRewriter &rewriter) {
+        Torch::ValueTensorType resultType;
+        Value input;
+        SmallVector<int64_t> axes;
+
+        if (binder.tensorOperand(input) ||
+            binder.s64IntegerArrayAttr(axes, "axes",
+                                       llvm::SmallVector<int64_t>({0, 2, 3})) ||
+            binder.tensorResultType(resultType)) {
+          return failure();
+        }
+        Location loc = binder.getLoc();
+        Value keepDim = rewriter.create<Torch::ConstantBoolOp>(loc, true);
+        Value unBiased = rewriter.create<Torch::ConstantBoolOp>(loc, false);
+        Value none = rewriter.create<Torch::ConstantNoneOp>(loc);
+
+        ArrayRef<int64_t> input_shape = resultType.getSizes();
+        SmallVector<int64_t> reduced_shape(input_shape);
+        for (int64_t i : axes) {
+          reduced_shape[i] = 1;
+        }
+
+        Torch::ValueTensorType meanOutTy = Torch::ValueTensorType::get(
+            resultType.getContext(), reduced_shape, resultType.getDtype());
+        SmallVector<Value> cstAxes;
+        for (int64_t i : axes) {
+          cstAxes.push_back(rewriter.create<Torch::ConstantIntOp>(
+              loc, rewriter.getI64IntegerAttr(i)));
+        }
+        Value axes_list = rewriter.create<Torch::PrimListConstructOp>(
+            loc,
+            Torch::ListType::get(Torch::IntType::get(binder.op->getContext())),
+            cstAxes);
+        Value mean = rewriter.create<Torch::AtenMeanDimOp>(
+            loc, meanOutTy, input, axes_list, keepDim, none);
+
+        Value variance = rewriter.create<Torch::AtenVarDimOp>(
+            loc, meanOutTy, input, axes_list, unBiased, keepDim);
+
+        Value cstOne = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(1));
+        Value cstEps = rewriter.create<Torch::ConstantFloatOp>(
+            loc, rewriter.getF64FloatAttr(1e-9));
+        Value positiveVariance = rewriter.create<Torch::AtenAddScalarOp>(
+            loc, meanOutTy, variance, cstEps, cstOne);
+
+        Value sqrt = rewriter.create<Torch::AtenSqrtOp>(loc, meanOutTy,
+                                                        positiveVariance);
+
+        Value subValue = rewriter.create<Torch::AtenSubTensorOp>(
+            loc, resultType, input, mean, cstOne);
+
+        Value meanVarNorm = rewriter.create<Torch::AtenDivTensorOp>(
+            loc, resultType, subValue, sqrt);
+
+        rewriter.replaceOp(binder.op, meanVarNorm);
+
+        return success();
+      });
   patterns.onOp(
       "Max", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
         Torch::ValueTensorType resultType;