A couple of extensions to MHA fusion

[gramalingam]

A couple of extensions to MHA fusion:

• One deals with variations in positions-ids. The challenge is to verify that the position-ids used in the two RotaryEmbedding are the same. In some models, they are the same value (by reference). In some models, there is some duplication of the code in computing the 2D position-id from 1D position-id. If we had a common-sub-expression identification/elimination, that would help. For now, just handling it in the pattern itself.

• The second deals with variations in how the last two axes of key are transposed. Some models reshape the input tensor to 3D and do the transpose, while some directly transpose a 4D tensor.

[codecov]

❌ 21 Tests Failed:

Tests completed	Failed	Passed	Skipped
14581	21	14560	2403

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

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0368_test_expand_dim_unchanged

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: 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_expand_dim_unchanged'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_expand_dim_unchanged' (e=No module named 'tests.onnx_backend_test_code.test_expand_dim_unchanged') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_expand_dim_unchanged.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_expand_dim_unchanged.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
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, INT64
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_expand_dim_unchanged(data: FLOAT[3,1], new_shape: INT64[2]) -> (FLOAT[3,4]):
E       expanded = opset13.Expand(data, new_shape)
E       return expanded

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0418_test_greater_bcast

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: 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_greater_bcast'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_greater_bcast' (e=No module named 'tests.onnx_backend_test_code.test_greater_bcast') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_greater_bcast.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_greater_bcast.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
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 BOOL, FLOAT
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_greater_bcast(x: FLOAT[3,4,5], y: FLOAT[5]) -> (BOOL[3,4,5]):
E       greater = opset13.Greater(x, y)
E       return greater

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0806_test_reduce_l1_do_not_keepdims_example_expanded

Stack Traces | 0.004s run time

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

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_reduce_l1_do_not_keepdims_example_expanded' (e=No module named 'tests.onnx_backend_test_code.test_reduce_l1_do_not_keepdims_example_expanded') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l1_do_not_keepdims_example_expanded.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l1_do_not_keepdims_example_expanded.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
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, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_reduce_l1_do_not_keepdims_example_expanded(data: FLOAT[3,2,2], axes: INT64[1]) -> (FLOAT[3,2]):
E       ReduceL1_test_reduce_l1_do_not_keepdims_example_expanded_function_data_abs = opset18.Abs(data)
E       reduced = opset18.ReduceSum(ReduceL1_test_reduce_l1_do_not_keepdims_example_expanded_function_data_abs, axes, keepdims=0)
E       return reduced

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