Move into onnx fusion folder

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

…nxscript/rewriter/layer_normalization.py …ipt/rewriter/onnx_fusions/_layer_norm.pyonnxscript/rewriter/layer_normalization.py renamed to onnxscript/rewriter/onnx_fusions/_layer_norm.py onnxscript/rewriter/layer_normalization.py renamed to onnxscript/rewriter/onnx_fusions/_layer_norm.py

@@ -34,6 +34,11 @@
 
 
 class LayerNormFusion(pattern.RewriteRuleClassBase):
+    def __init__(self, name: str, has_bias: bool):
+        super().__init__(name)
+        self._has_bias = has_bias
+        self._stash_dtype: int | None = None
+
     def pattern(self, op, x, scale, bias, epsilon, target_dtype):
         # Compute mean: Mean = ReduceMean(X, axes=normalized_axes)
         # TODO: support axes attribute too
@@ -43,8 +48,10 @@ def pattern(self, op, x, scale, bias, epsilon, target_dtype):
         deviation = op.Sub(x, mean)
 
         # Compute squared deviation: DD = Mul(D, D)
-        # TODO: support Pow (D, 2) as well
-        deviation_squared = op.Mul(deviation, deviation)
+        deviation_squared = pattern.OrValue([
+            op.Mul(deviation, deviation),
+            op.Pow(deviation, 2),
+        ])
 
         # Compute variance: Var = ReduceMean(DD, axes=normalized_axes)
         variance = op.ReduceMean(deviation_squared, [-1], keepdims=1)
@@ -56,24 +63,25 @@ def pattern(self, op, x, scale, bias, epsilon, target_dtype):
         std_dev = op.Sqrt(variance_plus_epsilon)
 
         # Compute reciprocal: InvStdDev = Reciprocal(StdDev)
-        # TODO: support Div(deviation, std_dev) as well?
-        inv_std_dev = op.Reciprocal(std_dev)
-
         # Normalize: Normalized = Mul(D, InvStdDev)
-        normalized = op.Mul(deviation, inv_std_dev)
+
+        inv_std_dev = op.Reciprocal(std_dev)
+        normalized = pattern.OrValue([
+            op.Mul(deviation, inv_std_dev),
+            op.Div(deviation, std_dev)
+        ])
 
         # Scale: NormalizedScaled = Mul(Normalized, Scale)
         normalized_scaled = op.Mul(normalized, scale)
 
         # Add bias (if present): Y = Add(NormalizedScaled, B)
-        if bias is not None:
+        
+        if self._has_bias:
             return op.Add(normalized_scaled, bias)
         else:
             return normalized_scaled
 
-    def check(
-        self, op, x, scale, bias, epsilon, compute_dtype, target_dtype, **_
-    ) -> pattern.MatchResult:  # type: ignore[name-defined]
+    def check(self, op, x, scale, bias, epsilon, target_dtype, **_) -> pattern.MatchResult:  # type: ignore[name-defined]
         """Check if the pattern matches conditions for use of LayerNormalization op."""
         check_result = pattern.MatchResult()
 

…ipt/rewriter/layer_normalization_test.py …ewriter/onnx_fusions/_layer_norm_test.pyonnxscript/rewriter/layer_normalization_test.py renamed to onnxscript/rewriter/onnx_fusions/_layer_norm_test.py onnxscript/rewriter/layer_normalization_test.py renamed to onnxscript/rewriter/onnx_fusions/_layer_norm_test.py

@@ -9,7 +9,7 @@
 import onnxscript.rewriter.testing
 from onnxscript import FLOAT, OnnxFunction, script
 from onnxscript import opset18 as op
-from onnxscript.rewriter.layer_normalization import fuse_layer_normalization
+from onnxscript.rewriter.onnx_fusions._layer_norm import fuse_layer_normalization
 
 
 @script()
@@ -86,13 +86,7 @@ def _test_layer_norm_with_bias(
 
 
 class LayerNormFusionTest(unittest.TestCase):
-    def _check(
-        self,
-        test_script: OnnxFunction,
-        expected_graph_len: int,
-        expected_op_type: str,
-        has_bias: bool = False,
-    ):
+    def _check(self, test_script: OnnxFunction):
         """Helper method to run a fusion test scenario."""
         model_proto = test_script.to_model_proto()
         # Create test inputs
@@ -112,11 +106,11 @@ def _check(
 
     def test_layer_norm_fusion_without_bias(self):
         """Test LayerNorm fusion without bias."""
-        self._check(_test_layer_norm_without_bias, 1, "LayerNormalization", has_bias=False)
+        self._check(_test_layer_norm_without_bias)
 
     def test_layer_norm_fusion_with_bias(self):
         """Test LayerNorm fusion with bias."""
-        self._check(_test_layer_norm_with_bias, 1, "LayerNormalization", has_bias=True)
+        self._check(_test_layer_norm_with_bias)
 
 
 if __name__ == "__main__":

onnxscript/rewriter/testing.py

@@ -11,7 +11,7 @@
 from onnxscript import ir
 
 
-def generate_random_inputs(self, model: onnx.ModelProto) -> dict[str, Any]:
+def generate_random_inputs(model: onnx.ModelProto) -> dict[str, Any]:
     feeds: dict[str, Any] = {}
     for input in model.graph.input:
         input_type = input.type.tensor_type