Merge pull request #284 from Xilinx/rogarcia.fix_layerNorm_recomposition

Fix LayerNorm recomposition by properly using provided epsilon value
Xilinx · Feb 5, 2025 · f4d1ac4 · f4d1ac4
2 parents a96118c + 5051f97
commit f4d1ac4
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 35 deletions.
diff --git a/src/Dialect/ONNX/Transforms/Recompose.cpp b/src/Dialect/ONNX/Transforms/Recompose.cpp
@@ -244,7 +244,9 @@ struct RecomposeLayerNormFromMulPattern : public OpRewritePattern<ONNXMulOp> {
         mlir::dyn_cast<ONNXConstantOp>(epsilon.getDefiningOp());
     if (!epsilonOp)
       return reportFailure("RMS epsilon needs to be a constant");
-    epsilonAttr = epsilonOp.getValueFloatAttr();
+    const auto epsilonValue = getScalarValue<double>(epsilonOp);
+    epsilonAttr =
+        FloatAttr::get(Float32Type::get(epsilonOp->getContext()), epsilonValue);
     // Check axes.
     if (!hasShapeAndRank(dd))
       return reportFailure("RMS need rank and shape for input dd");