Introduce layer-norm fusion

[gramalingam]

Introduce layer-norm fusion rules, along with a couple of test cases.

This is just the first version. TO DO:

• We need improved infrastructure for ONNX fusions to handle opset dependence. LayerNorm exists in ONNX since opset 17. For now the fusion rule exists, but it is not automatically called yet (but users can invoke it themselves).
• If users want to use opsets < 17, this could be done as an ORT fusion using ORT contrib op LayerNorm.

[codecov]

❌ 3 Tests Failed:

Tests completed	Failed	Passed	Skipped
15956	3	15953	2893

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0306_test_cumsum_1d_reverse_exclusive

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:132: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_cumsum_1d_reverse_exclusive'

The above exception was the direct cause of the following exception:
.nox\test_ort_nightly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:266: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:134: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_cumsum_1d_reverse_exclusive' (e=No module named 'tests.onnx_backend_test_code.test_cumsum_1d_reverse_exclusive') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_cumsum_1d_reverse_exclusive.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_cumsum_1d_reverse_exclusive.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy as np
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import DOUBLE, INT32
E   from onnxscript.onnx_opset import opset14
E   
E   @script()
E   def bck_test_cumsum_1d_reverse_exclusive(x: DOUBLE[5], axis: INT32) -> (DOUBLE[5]):
E       y = opset14.CumSum(x, axis, exclusive=1, reverse=1)
E       return y

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0543_test_layer_normalization_4d_axis_negative_3

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:132: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_layer_normalization_4d_axis_negative_3'

The above exception was the direct cause of the following exception:
.nox\test_ort_nightly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:266: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:134: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_layer_normalization_4d_axis_negative_3' (e=No module named 'tests.onnx_backend_test_code.test_layer_normalization_4d_axis_negative_3') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_4d_axis_negative_3.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_layer_normalization_4d_axis_negative_3.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy as np
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset17
E   
E   @script()
E   def bck_test_layer_normalization_4d_axis_negative_3(X: FLOAT[2,3,4,5], W: FLOAT[3,4,5], B: FLOAT[3,4,5]) -> (FLOAT[2,3,4,5], FLOAT[2,1,1,1], FLOAT[2,1,1,1]):
E       Y, Mean, InvStdDev = opset17.LayerNormalization(X, W, B, axis=-3)
E       return Y, Mean, InvStdDev

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1217_test_tfidfvectorizer_tf_batch_uniandbigrams_skip5

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:132: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_tfidfvectorizer_tf_batch_uniandbigrams_skip5'

The above exception was the direct cause of the following exception:
.nox\test_ort_nightly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:266: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:134: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_tfidfvectorizer_tf_batch_uniandbigrams_skip5' (e=No module named 'tests.onnx_backend_test_code.test_tfidfvectorizer_tf_batch_uniandbigrams_skip5') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_tfidfvectorizer_tf_batch_uniandbigrams_skip5.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_tfidfvectorizer_tf_batch_uniandbigrams_skip5.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy as np
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT32
E   from onnxscript.onnx_opset import opset9
E   
E   @script()
E   def bck_test_tfidfvectorizer_tf_batch_uniandbigrams_skip5(X: INT32[2,6]) -> (FLOAT[2,7]):
E       Y = opset9.TfIdfVectorizer(X, max_gram_length=2, max_skip_count=5, min_gram_length=1, mode='TF', ngram_counts=[0, 4], ngram_indexes=[0, 1, 2, 3, 4, 5, 6], pool_int64s=[2, 3, 5, 4, 5, 6, 7, 8, 6, 7])
E       return Y

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

Pull Request Overview

This PR introduces layer normalization fusion functionality to optimize ONNX models by fusing layer normalization patterns into the ONNX LayerNormalization operator. This is designed for ONNX opset 17+ where the LayerNormalization operator is available.

Key changes:

• Adds fusion rules to detect and fuse layer normalization patterns with and without bias
• Implements pattern matching for the mathematical sequence of operations that constitute layer normalization
• Includes comprehensive test cases to validate the fusion functionality

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
onnxscript/rewriter/testing.py	Adds utility functions for generating random inputs and enhances numerical equality testing to support dictionary-based inputs
onnxscript/rewriter/onnx_fusions/_layer_norm.py	Implements the core layer normalization fusion logic with pattern matching and rewrite rules
onnxscript/rewriter/onnx_fusions/_layer_norm_test.py	Provides test cases for layer normalization fusion with and without bias scenarios