[IR] Create a shape inference pass using onnx shape inference #2117

justinchuby · 2025-03-21T00:22:05Z

It handles large models by removing the initializers before sending the model to onnx shape inference.

Copilot

Pull Request Overview

This PR introduces a new shape inference pass that leverages onnx.shape_inference to automatically determine tensor shapes in the model. It adds the implementation for the pass, corresponding unit tests, and a minor cleanup in the model saving routine.

Reviewed Changes

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

File	Description
onnxscript/ir/passes/common/shape_inference.py	Implements the new shape inference pass using onnx.shape_inference.
onnxscript/ir/passes/common/shape_inference_test.py	Adds tests to validate that the pass correctly infers shapes and data types.
onnxscript/ir/_io.py	Updates the initializer tensor extraction for consistency in the model saving function.

Comments suppressed due to low confidence (1)

onnxscript/ir/passes/common/shape_inference_test.py:24

[nitpick] The node is created with an empty string as its name; consider providing a descriptive name to improve debugging and traceability.

add_node = ir.Node("", "Add", inputs=inputs)

codecov · 2025-03-21T00:26:17Z

❌ 4 Tests Failed:

Tests completed	Failed	Passed	Skipped
15548	4	15544	2756

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

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.008s run time

onnxscript/converter.py:467: 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:91: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:35: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1460: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1447: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:969: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1056: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:1000: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:553: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:832: in _translate_call_expr
    attrs = [
onnxscript/converter.py:833: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:517: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:469: 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_0397_test_ai_onnx_ml_tree_ensemble_set_membership

Stack Traces | 0.014s run time

onnxscript/converter.py:467: 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.11.../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:91: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:35: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1460: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1447: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:969: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1056: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:1000: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:553: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:832: in _translate_call_expr
    attrs = [
onnxscript/converter.py:833: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:517: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:469: 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.051s run time

onnxscript\converter.py:467: 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:91: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript\main.py:35: in script_check
    return convert.translate_function_def(f)
onnxscript\converter.py:1460: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript\converter.py:1447: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript\converter.py:969: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript\converter.py:1056: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript\converter.py:1000: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript\converter.py:553: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript\converter.py:832: in _translate_call_expr
    attrs = [
onnxscript\converter.py:833: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript\converter.py:517: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript\converter.py:469: 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/ir/passes/common/shape_inference.py

onnxscript/ir/passes/common/shape_inference_test.py

onnxscript/ir/passes/common/shape_inference.py

Copilot

Pull Request Overview

This PR introduces a new shape inference pass that temporarily removes large initializers to reduce model size for ONNX shape inference, and adds corresponding tests to ensure the pass functions as expected.

Implements a shape inference pass in onnxscript/ir/passes/common/shape_inference.py.
Adds tests to verify correct inference and the restoration of initializers in onnxscript/ir/passes/common/shape_inference_test.py.
Updates documentation in _pass_infra.py and refactors initializer value handling in _io.py to improve clarity.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
onnxscript/ir/passes/common/shape_inference_test.py	Added tests verifying shape and type inference on node outputs.
onnxscript/ir/passes/common/shape_inference.py	Implementation of the shape inference pass using onnx.shape_inference.
onnxscript/ir/passes/_pass_infra.py	Updated docstring to clarify the meaning of the "modified" attribute.
onnxscript/ir/_io.py	Simplified variable usage for initializing tensor backup.

Comments suppressed due to low confidence (1)

onnxscript/ir/passes/common/shape_inference_test.py:38

Consider adding a test case where an initializer's const_value exceeds _BIG_TENSOR_SIZE_LIMIT to validate that large initializers are temporarily removed and then properly restored after shape inference.

# Perform shape inference

onnxscript/ir/passes/common/shape_inference_test.py

Copilot

Pull Request Overview

This PR introduces a new shape inference pass leveraging onnx.shape_inference to support large models by temporarily removing oversized initializers.

Added the ShapeInferencePass class and an infer_shapes function in onnxscript/ir/passes/common/shape_inference.py
Introduced tests to verify the shape inference behavior and initializer restoration in onnxscript/ir/passes/common/shape_inference_test.py
Made minor documentation and code adjustments in onnxscript/ir/passes/_pass_infra.py and onnxscript/ir/_io.py to improve clarity and consistency

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
onnxscript/ir/passes/common/shape_inference.py	New shape inference pass implementation
onnxscript/ir/passes/common/shape_inference_test.py	Added tests for the shape inference pass
onnxscript/ir/passes/_pass_infra.py	Updated PassResult docstring
onnxscript/ir/_io.py	Refactored initializer value extraction for clarity

onnxscript/ir/passes/common/shape_inference_test.py

onnxscript/ir/passes/common/shape_inference.py

gramalingam

LGTM, thanks!

…oft#2117) It handles large models by removing the initializers before sending the model to onnx shape inference.

justinchuby added 3 commits March 20, 2025 17:13

[IR] Create a shape inference pass using onnx shape inference

5b3de02

lint

f3e2f5f

test

fce00f1

justinchuby requested a review from Copilot March 21, 2025 00:22

github-project-automation bot added this to ONNX Script Review Board Mar 21, 2025

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

Copilot AI reviewed Mar 21, 2025

View reviewed changes

justinchuby added the topic: passes label Mar 21, 2025

justinchuby requested a review from gramalingam March 21, 2025 00:22

justinchuby mentioned this pull request Mar 21, 2025

Cleanup ort transformer fusions #2115

Merged

warning

f04af76

justinchuby added 2 commits March 20, 2025 17:26

strict

b065bfb

config

7292f25

github-advanced-security bot found potential problems Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Fixed Show fixed Hide fixed

justinchuby added 3 commits March 20, 2025 17:40

restore

2bad82a

docs

0eb0761

shape

999d9f1

github-advanced-security bot found potential problems Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Fixed Show fixed Hide fixed

onnxscript/ir/passes/common/shape_inference.py Fixed Show fixed Hide fixed

__all__

fccd8f4

justinchuby commented Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Outdated Show resolved Hide resolved

justinchuby commented Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Outdated Show resolved Hide resolved

Apply suggestions from code review

24087df

github-advanced-security bot found potential problems Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Fixed Show fixed Hide fixed

onnxscript/ir/passes/common/shape_inference.py Fixed Show fixed Hide fixed

justinchuby commented Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Outdated Show resolved Hide resolved

Update onnxscript/ir/passes/common/shape_inference.py

7385405

gramalingam reviewed Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference_test.py Show resolved Hide resolved

gramalingam reviewed Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Outdated Show resolved Hide resolved

gramalingam reviewed Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Show resolved Hide resolved

Address comments

9ac3fe9

Merge branch 'main' into justinchu/shape-inference

1251e94

justinchuby requested review from Copilot and gramalingam March 21, 2025 23:17

Copilot AI reviewed Mar 21, 2025

View reviewed changes

shubhambhokare1 reviewed Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference_test.py Outdated Show resolved Hide resolved

more tests

2eed545

justinchuby requested review from Copilot and shubhambhokare1 March 21, 2025 23:51

Copilot AI reviewed Mar 21, 2025

View reviewed changes

github-advanced-security bot found potential problems Mar 21, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference_test.py Fixed Show fixed Hide fixed

test

fce63b3

justinchuby marked this pull request as draft March 22, 2025 14:59

gramalingam reviewed Mar 24, 2025

View reviewed changes

onnxscript/ir/passes/common/shape_inference.py Show resolved Hide resolved

gramalingam approved these changes Mar 24, 2025

View reviewed changes

justinchuby marked this pull request as ready for review March 24, 2025 17:30

justinchuby enabled auto-merge (squash) March 24, 2025 17:32

justinchuby disabled auto-merge March 24, 2025 17:32

Update modified

d69fb2d

justinchuby enabled auto-merge (squash) March 24, 2025 17:38

justinchuby merged commit 89b7a05 into main Mar 24, 2025
22 of 27 checks passed

justinchuby deleted the justinchu/shape-inference branch March 24, 2025 17:54

github-project-automation bot moved this from Todo to Done in ONNX Script Review Board Mar 24, 2025

bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025

[IR] Create a shape inference pass using onnx shape inference (micros…

f90c2d8

…oft#2117) It handles large models by removing the initializers before sending the model to onnx shape inference.

bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025

[IR] Create a shape inference pass using onnx shape inference (micros…

9209aff

…oft#2117) It handles large models by removing the initializers before sending the model to onnx shape inference.

bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025

[IR] Create a shape inference pass using onnx shape inference (micros…

175240b

…oft#2117) It handles large models by removing the initializers before sending the model to onnx shape inference.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[IR] Create a shape inference pass using onnx shape inference #2117

[IR] Create a shape inference pass using onnx shape inference #2117

Uh oh!

justinchuby commented Mar 21, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

codecov bot commented Mar 21, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

gramalingam left a comment

Uh oh!

Uh oh!

Uh oh!

[IR] Create a shape inference pass using onnx shape inference #2117

[IR] Create a shape inference pass using onnx shape inference #2117

Uh oh!

Conversation

justinchuby commented Mar 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

codecov bot commented Mar 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ 4 Tests Failed:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

gramalingam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Mar 21, 2025 •

edited

Loading

codecov bot commented Mar 21, 2025 •

edited

Loading