Minor fixes

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

Author: gramalingam

onnxscript/rewriter/layer_normalization.py

@@ -9,7 +9,7 @@
 """
 Layer Normalization fusion optimization.
 
-This module contains rewrite rules for fusing Layer Normalization patterns into the 
+This module contains rewrite rules for fusing Layer Normalization patterns into the
 ONNX LayerNormalization operator.
 
 Layer Normalization performs normalization over the last D dimensions as specified by the axis.
@@ -34,37 +34,37 @@
 
 
 class LayerNormFusion(pattern.RewriteRuleClassBase):
-    def pattern(self, op, x, scale, bias, epsilon, target_dtype):       
+    def pattern(self, op, x, scale, bias, epsilon, target_dtype):
         # Compute mean: Mean = ReduceMean(X, axes=normalized_axes)
         # TODO: support axes attribute too
         mean = op.ReduceMean(x, [-1], keepdims=1)
-        
+
         # Compute deviation: D = Sub(X, Mean)
         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)
-        
+
         # Compute variance: Var = ReduceMean(DD, axes=normalized_axes)
         variance = op.ReduceMean(deviation_squared, [-1], keepdims=1)
-        
+
         # Add epsilon: VarEps = Add(Var, epsilon)
         variance_plus_epsilon = op.Add(variance, epsilon)
-        
+
         # Compute standard deviation: StdDev = Sqrt(VarEps)
         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)
-       
+
         # Scale: NormalizedScaled = Mul(Normalized, Scale)
         normalized_scaled = op.Mul(normalized, scale)
-        
+
         # Add bias (if present): Y = Add(NormalizedScaled, B)
         if bias is not None:
             return op.Add(normalized_scaled, bias)
@@ -76,17 +76,17 @@ def check(
     ) -> pattern.MatchResult:  # type: ignore[name-defined]
         """Check if the pattern matches conditions for use of LayerNormalization op."""
         check_result = pattern.MatchResult()
-        
+
         # epsilon must be a scalar
         epsilon_value = _ir_utils.get_singleton_value(epsilon)
         if not isinstance(epsilon_value, float):  # TODO: support other types
             return check_result.fail("Epsilon is not a float value.", epsilon)
-        
+
         if x.dtype not in fp_float_types:
             return check_result.fail("Input is not a float type.", x)
-           
+
         self._stash_dtype = x.dtype
-        
+
         return check_result
 
     def rewrite(self, op, x, scale, bias, epsilon, **_):

onnxscript/rewriter/layer_normalization_test.py

@@ -3,120 +3,112 @@
 
 import unittest
 
-import numpy as np
 import onnx_ir as ir
-import parameterized
 
 import onnxscript
-import onnxscript.rewriter.ort_fusions._test_utils as test_utils
+import onnxscript.rewriter.testing
 from onnxscript import FLOAT, OnnxFunction, script
-from onnxscript import opset17 as op
-from onnxscript.optimizer import optimize, remove_unused_nodes
+from onnxscript import opset18 as op
 from onnxscript.rewriter.layer_normalization import fuse_layer_normalization
-import onnxscript.rewriter.testing
 
 
 @script()
 def _test_layer_norm_without_bias(x: FLOAT[2, 4, 8], scale: FLOAT[8]) -> FLOAT[2, 4, 8]:
     """LayerNorm pattern without bias."""
     # Compute mean: Mean = ReduceMean(X, axes=normalized_axes)
     mean = op.ReduceMean(x, [-1], keepdims=1)
-    
+
     # Compute deviation: D = Sub(X, Mean)
     deviation = op.Sub(x, mean)
-    
+
     # Compute squared deviation: DD = Mul(D, D)
     deviation_squared = op.Mul(deviation, deviation)
-    
+
     # Compute variance: Var = ReduceMean(DD, axes=normalized_axes)
     variance = op.ReduceMean(deviation_squared, [-1], keepdims=1)
-    
+
     # Add epsilon: VarEps = Add(Var, epsilon)
     epsilon = op.Constant(value_float=1e-5)
     variance_plus_epsilon = op.Add(variance, epsilon)
-    
+
     # Compute standard deviation: StdDev = Sqrt(VarEps)
     std_dev = op.Sqrt(variance_plus_epsilon)
-    
+
     # Compute reciprocal: InvStdDev = Reciprocal(StdDev)
     inv_std_dev = op.Reciprocal(std_dev)
-    
+
     # Normalize: Normalized = Mul(D, InvStdDev)
     normalized = op.Mul(deviation, inv_std_dev)
-    
+
     # Scale: NormalizedScaled = Mul(Normalized, Scale)
     normalized_scaled = op.Mul(normalized, scale)
-    
+
     return normalized_scaled
 
 
 @script()
-def _test_layer_norm_with_bias(x: FLOAT[2, 4, 8], scale: FLOAT[8], bias: FLOAT[8]) -> FLOAT[2, 4, 8]:
+def _test_layer_norm_with_bias(
+    x: FLOAT[2, 4, 8], scale: FLOAT[8], bias: FLOAT[8]
+) -> FLOAT[2, 4, 8]:
     """LayerNorm pattern with bias."""
     # Compute mean: Mean = ReduceMean(X, axes=normalized_axes)
     mean = op.ReduceMean(x, [-1], keepdims=1)
-    
+
     # Compute deviation: D = Sub(X, Mean)
     deviation = op.Sub(x, mean)
-    
+
     # Compute squared deviation: DD = Mul(D, D)
     deviation_squared = op.Mul(deviation, deviation)
-    
+
     # Compute variance: Var = ReduceMean(DD, axes=normalized_axes)
     variance = op.ReduceMean(deviation_squared, [-1], keepdims=1)
-    
+
     # Add epsilon: VarEps = Add(Var, epsilon)
     epsilon = op.Constant(value_float=1e-5)
     variance_plus_epsilon = op.Add(variance, epsilon)
-    
+
     # Compute standard deviation: StdDev = Sqrt(VarEps)
     std_dev = op.Sqrt(variance_plus_epsilon)
-    
+
     # Compute reciprocal: InvStdDev = Reciprocal(StdDev)
     inv_std_dev = op.Reciprocal(std_dev)
-    
+
     # Normalize: Normalized = Mul(D, InvStdDev)
     normalized = op.Mul(deviation, inv_std_dev)
-    
+
     # Scale: NormalizedScaled = Mul(Normalized, Scale)
     normalized_scaled = op.Mul(normalized, scale)
-    
+
     # Add bias: Y = Add(NormalizedScaled, B)
     result = op.Add(normalized_scaled, bias)
-    
+
     return result
 
 
 class LayerNormFusionTest(unittest.TestCase):
     def _check(
         self,
-        test_data_constructor: OnnxFunction,
+        test_script: OnnxFunction,
         expected_graph_len: int,
         expected_op_type: str,
         has_bias: bool = False,
     ):
         """Helper method to run a fusion test scenario."""
-        model_proto = test_data_constructor.to_model_proto()
+        model_proto = test_script.to_model_proto()
         # Create test inputs
-        input_data = onnxscript.rewriter.testing.generate_random_inputs(model)
+        input_data = onnxscript.rewriter.testing.generate_random_inputs(model_proto)
 
         model = ir.serde.deserialize_model(model_proto)
         fuse_layer_normalization(model)
 
-        # Run original model
-        original_output = test_utils.ort_run("Original", model, input_data)
-
-        # Apply fusion
-        fuse_layer_normalization(model)
-        remove_unused_nodes(model)
+        # Check that a LayerNormalization node was created
+        self.assertIn("LayerNormalization", [n.op_type for n in model.graph])
 
-        # Verify fusion occurred
-        self.assertEqual(len(model.graph), expected_graph_len)
-        self.assertEqual(model.graph.node(0).op_type, expected_op_type)
+        fused_model_proto = ir.serde.serialize_model(model)
 
-        # Run optimized model and verify outputs match
-        optimized_output = test_utils.ort_run("Optimized", model, input_data)
-        test_utils.assert_allclose(original_output, optimized_output, rtol=1e-4, atol=1e-4)
+        onnxscript.rewriter.testing.assert_numerically_equal(
+            model_proto, fused_model_proto, input_data
+        )
 
     def test_layer_norm_fusion_without_bias(self):
         """Test LayerNorm fusion without bias."""
@@ -128,4 +120,4 @@ def test_layer_norm_fusion_with_bias(self):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()

onnxscript/rewriter/testing.py

@@ -10,12 +10,13 @@
 
 from onnxscript import ir
 
+
 def generate_random_inputs(self, model: onnx.ModelProto) -> dict[str, Any]:
     feeds: dict[str, Any] = {}
     for input in model.graph.input:
         input_type = input.type.tensor_type
         shape = tuple(input_type.shape.dim)
-        if not all(hasattr(d, 'dim_value') for d in shape):
+        if not all(hasattr(d, "dim_value") for d in shape):
             raise ValueError(f"Input {input.name} has dynamic shape dimensions.")
         shape = tuple(d.dim_value for d in shape)
         if input_type.elem_type == onnx.TensorProto.FLOAT:
@@ -28,11 +29,10 @@ def generate_random_inputs(self, model: onnx.ModelProto) -> dict[str, Any]:
     return feeds
 
 
-
 def assert_numerically_equal(
     original_model_proto: onnx.ModelProto | ir.Model,
     rewritten_model_proto: onnx.ModelProto | ir.Model,
-    args: tuple[Any, ...],
+    args: tuple[Any, ...] | dict[str, Any],
     ort_optimization_level: ort.GraphOptimizationLevel = ort.GraphOptimizationLevel.ORT_ENABLE_ALL,
     rtol: float = 1,
     atol: float = 1e-3,
@@ -53,9 +53,17 @@ def assert_numerically_equal(
     if isinstance(rewritten_model_proto, ir.Model):
         rewritten_model_proto = ir.serde.serialize_model(rewritten_model_proto)
 
-    original_proto_ort_inputs = {
-        k.name: v for k, v in zip(original_model_proto.graph.input, args)
-    }
+    if isinstance(args, dict):
+        original_proto_ort_inputs = args
+        the_rewritten_proto_ort_inputs = args
+    else:
+        original_proto_ort_inputs = {
+            k.name: v for k, v in zip(original_model_proto.graph.input, args)
+        }
+        the_rewritten_proto_ort_inputs = {
+            k.name: v for k, v in zip(rewritten_model_proto.graph.input, args)
+        }
+
     original_proto_ort_inference_session = _ort_session_initializer(
         original_model_proto.SerializeToString(), ort_optimization_level
     )
@@ -65,9 +73,6 @@ def assert_numerically_equal(
         None, original_proto_ort_inputs, run_options=run_options
     )
 
-    the_rewritten_proto_ort_inputs = {
-        k.name: v for k, v in zip(rewritten_model_proto.graph.input, args)
-    }
     the_rewritten_proto_ort_inference_session = _ort_session_initializer(
         rewritten_model_proto.SerializeToString(), ort_optimization_level
     )