Remove the RefAttr class #2328

justinchuby · 2025-05-22T15:19:11Z

Rational

We defined the class RefAttr in the IR to represent reference attributes in ONNX. Node attributes can be Attr and RefAttr. However, since most of the time we are working with concrete attributes, the union of types creates a typing situation where we always need to assert the types before taking the values, even if we know a RefAttr cannot exist (outside of a function definition).

This additionally matches the definition of AttributeProto in ONNX.

Change

This change merged the two classes, and instead defines a is_ref() method for users to check the reference attribute.

The change is BC breaking for usage like isinstance(attr, ir.RefAttr). Fortunately all such usages exist in this code base and not in PyTorch, so we are safe to complete the change.

onnxscript/ir/_protocols.py

@@ -422,6 +423,8 @@
    value: Any
    doc_string: str | None

+    def is_ref(self) -> Literal[False]: ...


To fix the issue, replace the ... placeholder in the is_ref method with a concrete implementation that returns False. This aligns the method's behavior with its type annotation (Literal[False]) and ensures that the protocol is clear and unambiguous for implementers.

onnxscript/ir/_protocols.py

@@ -441,6 +444,8 @@
    type: _enums.AttributeType
    doc_string: str | None

+    def is_ref(self) -> Literal[True]: ...


To fix the issue, the is_ref method in the ReferenceAttributeProtocol class should be updated to explicitly return True as a Literal[True]. This ensures that the method has a concrete implementation that matches its type annotation and intended behavior. The change should be made directly in the ReferenceAttributeProtocol class definition.

codecov · 2025-05-22T15:21:54Z

❌ 3 Tests Failed:

Tests completed	Failed	Passed	Skipped
16177	3	16174	1703

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

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.016s 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                                                                                                                                                                                             ^

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0815_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:
..../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.025s 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                                                                                                                                                                                             ^

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 removes the separate RefAttr class by merging its functionality into the existing Attr class and introduces an is_ref() method for reference checks.
Key changes:

Consolidated attribute classes and removed RefAttr across the IR.
Updated all type annotations and isinstance(attr, RefAttr) checks to use Attr and attr.is_ref().
Adjusted serialization/deserialization and protocol definitions to handle reference attributes via is_ref().

Reviewed Changes

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

Show a summary per file

File	Description
onnxscript/version_converter/_version_converter.py	Changed `visit_attribute` signature to `ir.Attr`.
onnxscript/rewriter/llama_rule_sets.py	Replaced `isinstance(..., RefAttr)` with `is_ref()`.
onnxscript/rewriter/_rewrite_rule.py	Simplified `check` signatures; updated `copy_attr_value`.
onnxscript/rewriter/_pattern_ir.py	Narrowed `AttrPattern` to `ir.Attr` only.
onnxscript/optimizer/_constant_folding.py	Changed `visit_attribute` signature to `ir.Attr`.
onnxscript/ir/serde.py	Split serialization paths using `is_ref()`.
onnxscript/ir/passes/common/inliner.py	Updated `clone_attr` and `attr_map` types.
onnxscript/ir/external_data.py	Replaced `isinstance(..., RefAttr)` with `is_ref()`.
onnxscript/ir/_tape.py	Updated attribute sequences to `ir.Attr`.
onnxscript/ir/_protocols.py	Added `Literal` import and `is_ref()` to protocols.
onnxscript/ir/_core.py	Removed `RefAttr` class; merged into `Attr`.
onnxscript/ir/_convenience/_constructors.py	Updated attribute sequences to `ir.Attr`.
onnxscript/ir/_convenience/init.py	Removed `RefAttr` import and tightened types.

Comments suppressed due to low confidence (3)

onnxscript/version_converter/_version_converter.py:265

Reference attributes are no longer skipped here—since all attributes are now Attr, isinstance(attr, ir.Attr) will also match refs. Add a guard to skip ref attrs, e.g., if not attr.is_ref(): before processing.

if isinstance(attr, ir.Attr):

onnxscript/optimizer/_constant_folding.py:1063

Same as above, reference attributes will match ir.Attr and may cause visit_graph on a ref. Add not attr.is_ref() to skip processing reference attrs.

if isinstance(attr, ir.Attr):

onnxscript/rewriter/_rewrite_rule.py:371

The exception message still references the removed RefAttr class. Update it to something like "Reference attributes are not supported." for clarity.

raise NotImplementedError("RefAttr not supported.")

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

justinchuby added 2 commits May 22, 2025 08:10

Remove refattr usage

66cd9cf

Format

36bbd58

justinchuby added the topic: api label May 22, 2025

github-project-automation bot added this to ONNX Script Review Board May 22, 2025

justinchuby added topic: breaking changes module: IR Intermediate representation labels May 22, 2025

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

github-advanced-security bot found potential problems May 22, 2025

View reviewed changes

justinchuby mentioned this pull request May 22, 2025

Fix fused matmul check/rewrite functions #2322

Closed

Fix

6a7964a

justinchuby requested review from Copilot May 22, 2025 15:26

Copilot AI reviewed May 22, 2025

View reviewed changes

setter

dc78ea1

justinchuby added this to the 0.2.7 milestone May 22, 2025

justinchuby added 2 commits May 22, 2025 08:59

fix isinstance

b729215

graphs

0d355e5

justinchuby requested a review from Copilot May 22, 2025 16:45

Copilot AI reviewed May 22, 2025

View reviewed changes

xadupre approved these changes May 22, 2025

View reviewed changes

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

justinchuby merged commit f46004e into main May 22, 2025
26 of 29 checks passed

justinchuby deleted the justinchu/combine-ref branch May 22, 2025 17:43

justinchuby mentioned this pull request May 21, 2025

Migrate PRs onnx/ir-py#21

Closed

9 tasks

justinchuby modified the milestones: 0.2.7, 0.2.6 May 22, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Remove the RefAttr class #2328

Remove the RefAttr class #2328

Uh oh!

justinchuby commented May 22, 2025 •

edited

Loading

Uh oh!

Check notice

Copilot Autofix

Check notice

Copilot Autofix

codecov bot commented May 22, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Remove the RefAttr class #2328

Remove the RefAttr class #2328

Uh oh!

Conversation

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

Uh oh!

Rational

Change

Uh oh!

Check notice

Copilot Autofix

Check notice

Copilot Autofix

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

Uh oh!

❌ 3 Tests Failed:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

justinchuby commented May 22, 2025 •

edited

Loading

codecov bot commented May 22, 2025 •

edited

Loading