[rewriter] Transpose rule #2255

justinchuby · 2025-04-30T18:15:28Z

Create a rule to absorb a Transpose node into the initializer (fold constant). This rule is in place to handle initializers of any size. So users do not need to change the size limit of the optimizer when constant folding.

TODO

tests

Fix #2158

codecov · 2025-04-30T18:19:38Z

❌ 5 Tests Failed:

Tests completed	Failed	Passed	Skipped
6097	5	6092	3057

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

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0981_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.009s run time

onnxscript/converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../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:137: in extract_functions
    mod = importlib.import_module(import_name)
.../hostedtoolcache/Python/3.11.12.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:94: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:38: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:825: in _translate_call_expr
    attrs = [
onnxscript/converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0125_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.01s run time

onnxscript/converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../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:137: in extract_functions
    mod = importlib.import_module(import_name)
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:94: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:38: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:825: in _translate_call_expr
    attrs = [
onnxscript/converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0026_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.019s run time

onnxscript\converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

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: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: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)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
.nox\test_ort_nightly\Lib\site-packages\_pytest\assertion\rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests\onnx_backend_test_code\test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript\main.py:94: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript\main.py:38: in script_check
    return convert.translate_function_def(f)
onnxscript\converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript\converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript\converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript\converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript\converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript\converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript\converter.py:825: in _translate_call_expr
    attrs = [
onnxscript\converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript\converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript\converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

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

onnxscript/rewriter/transpose_initializer.py

Copilot

Pull Request Overview

This PR adds a new rewrite rule to collapse Transpose nodes into initializers for constant folding, which removes the need to modify the optimizer’s size limit when folding constants.

Introduces the TransposeInitializer rule to handle initializer folding.
Implements the rewriting and checking logic to ensure the initializer meets usage requirements.
Registers the new rule in the rewriter’s initializer.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxscript/rewriter/transpose_initializer.py	Adds a new rewrite rule for folding Transpose nodes into initializers.
onnxscript/rewriter/init.py	Registers the new transpose_initializer rule in the rule set.

Comments suppressed due to low confidence (1)

onnxscript/rewriter/transpose_initializer.py:28

[nitpick] Consider renaming 'original_transpose' to 'transpose_node' to more clearly indicate that it represents the Transpose node consumer of the initializer.

original_transpose = initializer.consumers()[0]

onnxscript/rewriter/transpose_initializer.py

justinchuby · 2025-04-30T19:16:53Z

I think this is less efficient than an IR pass. Because it will match all transpose nodes, but we only need to look at users of the initializers. Thoughts? @titaiwangms @shubhambhokare1 @gramalingam

justinchuby · 2025-04-30T19:19:17Z

But this is op-specific. I don't think it fits in the common passes. So I prefer that it is a rewrite rule.

onnxscript/rewriter/transpose_initializer.py

Co-authored-by: Shubham Bhokare <[email protected]>

titaiwangms · 2025-05-01T16:26:25Z

I think this is less efficient than an IR pass. Because it will match all transpose nodes, but we only need to look at users of the initializers. Thoughts? @titaiwangms @shubhambhokare1 @gramalingam

With the check inside the rule, I assume it should not be too bad?

titaiwangms · 2025-05-01T16:28:14Z

In the original issue, there are also transpose -> transpose and transpose -> MatMul/Gemm pattern matching. Should we have those in this rule as well? If we are going to go through Transpose nods anyway (I forget if we can check whether multiple rules are applicable within the same traverse.).

justinchuby · 2025-05-01T18:11:30Z

I think we should have different rules since they will match different nodes

shubhambhokare1 · 2025-05-07T15:59:42Z

Could we add

onnxscript/onnxscript/rewriter/generic_pattern_test.py

Lines 530 to 573 in 33f31ca

    
           def test_transpose_transpose_onnxscript(self): 
        
               # TODO(rama): Attribute-parameters not yet supported in multi-output matching. 
        
               # def transpose_transpose_pattern(op, X, perm0, perm1): 
        
               #     xt = op.Transpose(X, perm=perm0) 
        
               #     Y = op.Transpose(xt, perm=perm1) 
        
               #     return Y 
        
               def transpose_transpose_pattern(op, X): 
        
                   XT = op.Transpose(X, _outputs=["XT"]) 
        
                   Y = op.Transpose(XT, _outputs=["Y"]) 
        
                   return Y 
        
               def transpose_transpose_mapping(perm0, perm1): 
        
                   new_perm = [0 for p in perm0] 
        
                   for i, p in enumerate(perm1): 
        
                       new_perm[i] = perm0[p] 
        
                   # replace by return [perm0[p] for p in perm1] ? 
        
                   return new_perm 
        
               def transpose_transpose_check(op, **_) -> bool: 
        
                   return True 
        
               def transpose_transpose_apply_pattern(op, X, XT: ir.Value, Y, **_): 
        
                   perm0 = XT.producer().attributes.get("perm") 
        
                   if perm0 is not None: 
        
                       perm0 = perm0.value  # TODO(rama): handle RefAttr 
        
                   perm1 = Y.producer().attributes.get("perm") 
        
                   if perm1 is not None: 
        
                       perm1 = perm1.value  # TODO(rama): handle RefAttr 
        
                   if perm0 is None and perm1 is None: 
        
                       return op.Identity(X) 
        
                   if perm0 is None: 
        
                       perm0 = range(len(perm1) - 1, -1, -1) 
        
                   if perm1 is None: 
        
                       perm1 = range(len(perm0) - 1, -1, -1) 
        
                   composed_perm = transpose_transpose_mapping(perm0, perm1) 
        
                   return op.Transpose(X, perm=composed_perm) 
        
               rule = pattern.RewriteRule( 
        
                   transpose_transpose_pattern, 
        
                   transpose_transpose_apply_pattern, 
        
                   transpose_transpose_check, 
        
                   self.matcher_algo, 
        
                   verbose=0,

as part of this PR to upstream all the fusions done in https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/fusion_constant_fold.py

justinchuby · 2025-05-07T22:13:57Z

@shubhambhokare1 I see it's already in

onnxscript/onnxscript/rewriter/llama_rule_sets.py

Line 318 in 605e06e

transpose_transpose_rule,

, was it not applied?

shubhambhokare1 · 2025-05-09T18:03:13Z

@shubhambhokare1 I see it's already in

onnxscript/onnxscript/rewriter/llama_rule_sets.py

Line 318 in 605e06e

transpose_transpose_rule,

, was it not applied?

Ah I see, however the one in https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/fusion_constant_fold.py seems to be doing,
inp -> Transpose -> Transpose -> out
to
inp -> out if transpose perm values are equal

Whereas the current rewrite rule only replace two transposes with an equivalent single transpose

justinchuby · 2025-05-09T23:33:39Z

@shubhambhokare1 I see it's already in

onnxscript/onnxscript/rewriter/llama_rule_sets.py

Line 318 in 605e06e

transpose_transpose_rule,

, was it not applied?

Ah I see, however the one in microsoft/onnxruntime@main/onnxruntime/python/tools/transformers/fusion_constant_fold.py seems to be doing, inp -> Transpose -> Transpose -> out to inp -> out if transpose perm values are equal

Whereas the current rewrite rule only replace two transposes with an equivalent single transpose

Sounds good. I will update

justinchuby · 2025-05-10T00:13:59Z

@shubhambhokare1 does this one

onnxscript/onnxscript/rewriter/llama_rule_sets.py

Line 216 in ac87a1c

return op.Identity(x)

or this one

onnxscript/onnxscript/rewriter/llama_rule_sets.py

Line 257 in ac87a1c

if first == last:

do it?

justinchuby added 2 commits April 30, 2025 11:13

[rewriter] Transpose rule

0d8bece

WIP

43b9f26

github-project-automation bot added this to ONNX Script Review Board Apr 30, 2025

github-project-automation bot moved this to Todo in ONNX Script Review Board Apr 30, 2025

github-advanced-security bot found potential problems Apr 30, 2025

View reviewed changes

justinchuby added 3 commits April 30, 2025 11:54

Implement the rule

c9e93e1

init

37232d3

attr

0c47226

justinchuby marked this pull request as ready for review April 30, 2025 18:57

justinchuby added 2 commits April 30, 2025 11:59

new_name

f2d2ccf

Copyright

278e977

justinchuby requested a review from Copilot April 30, 2025 19:01

justinchuby added the module: rewriter label Apr 30, 2025

justinchuby requested review from titaiwangms and shubhambhokare1 April 30, 2025 19:02

Copilot AI reviewed Apr 30, 2025

View reviewed changes

onnxscript/rewriter/transpose_initializer.py Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Apr 30, 2025

View reviewed changes

onnxscript/rewriter/transpose_initializer.py Fixed Show fixed Hide fixed

justinchuby added 2 commits April 30, 2025 12:11

typing

d028482

todo

3c2369c

justinchuby marked this pull request as draft April 30, 2025 19:16

justinchuby marked this pull request as ready for review April 30, 2025 19:18

titaiwangms requested a review from gramalingam April 30, 2025 20:01

shubhambhokare1 reviewed Apr 30, 2025

View reviewed changes

onnxscript/rewriter/transpose_initializer.py Outdated Show resolved Hide resolved

Update onnxscript/rewriter/transpose_initializer.py

f271a7c

Co-authored-by: Shubham Bhokare <[email protected]>

justinchuby added 2 commits May 9, 2025 08:37

Merge branch 'main' into justinchu/transpose-rule

b39d763

update

19aaedb

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[rewriter] Transpose rule #2255

[rewriter] Transpose rule #2255

justinchuby commented Apr 30, 2025 •

edited

Loading

codecov bot commented Apr 30, 2025 •

edited

Loading

Copilot AI left a comment

justinchuby commented Apr 30, 2025

justinchuby commented Apr 30, 2025

titaiwangms commented May 1, 2025

titaiwangms commented May 1, 2025 •

edited

Loading

justinchuby commented May 1, 2025

shubhambhokare1 commented May 7, 2025

justinchuby commented May 7, 2025

shubhambhokare1 commented May 9, 2025

justinchuby commented May 9, 2025

justinchuby commented May 10, 2025

[rewriter] Transpose rule #2255

Are you sure you want to change the base?

[rewriter] Transpose rule #2255

Conversation

justinchuby commented Apr 30, 2025 • edited Loading

codecov bot commented Apr 30, 2025 • edited Loading

❌ 5 Tests Failed:

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

justinchuby commented Apr 30, 2025

justinchuby commented Apr 30, 2025

titaiwangms commented May 1, 2025

titaiwangms commented May 1, 2025 • edited Loading

justinchuby commented May 1, 2025

shubhambhokare1 commented May 7, 2025

justinchuby commented May 7, 2025

shubhambhokare1 commented May 9, 2025

justinchuby commented May 9, 2025

justinchuby commented May 10, 2025

justinchuby commented Apr 30, 2025 •

edited

Loading

codecov bot commented Apr 30, 2025 •

edited

Loading

titaiwangms commented May 1, 2025 •

edited

Loading