From 919763ebb15c4f3668bb5ed4b4c982a2d230caa2 Mon Sep 17 00:00:00 2001
From: "Xiaomei.Shang" <Xiaomei.Shang@synaptics.com>
Date: Wed, 22 Apr 2026 09:16:48 +0000
Subject: [PATCH] batchNorm and softmax

---
 .../TorchOnnxToTorch/DefaultDomainAtoF.cpp    |  2 +-
 .../TorchOnnxToTorch/DefaultDomainQtoZ.cpp    | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp
index eff45ee3d98a..6d6431199629 100644
--- a/lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp
+++ b/lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp
@@ -342,7 +342,7 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
                   return success();
                 });
   patterns.onOp(
-      "BatchNormalization", 15,
+      "BatchNormalization", 7,
       [](OpBinder binder, ConversionPatternRewriter &rewriter) {
         Torch::ValueTensorType resultType;
         Value input, weight, bias, inputMean, inputVar;
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
index 0de8ffa5d5d8..e7512e883794 100644
--- a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
+++ b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
@@ -1273,6 +1273,35 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
         rewriter.replaceOp(binder.op, result);
         return success();
       });
+  patterns.onOp(
+      "Softmax", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
+        Torch::ValueTensorType resultType;
+        Value input;
+        int64_t axis;
+        if (binder.tensorOperand(input) ||
+            binder.s64IntegerAttr(axis, "axis", -1) ||
+            binder.tensorResultType(resultType))
+          return failure();
+
+        auto inputType = cast<Torch::ValueTensorType>(input.getType());
+        if (!inputType.hasSizes()) {
+          return rewriter.notifyMatchFailure(
+              binder.op, "unimplemented: expected input to have sizes");
+        }
+
+        int64_t inputRank = inputType.getSizes().size();
+        // Handle negative axis
+        if (axis < 0)
+          axis = inputRank + axis;
+
+        Value constAxis = rewriter.create<Torch::ConstantIntOp>(
+            binder.getLoc(), rewriter.getI64IntegerAttr(axis));
+        Value dtypeNone = rewriter.create<Torch::ConstantNoneOp>(binder.getLoc());
+
+        rewriter.replaceOpWithNewOp<Torch::AtenSoftmaxIntOp>(
+            binder.op, resultType, input, constAxis, dtypeNone);
+        return success();
+      });
   patterns.onOp(
       "Softmax", 13, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
         Torch::ValueTensorType resultType;