[IR] Record owning graph for input/output/initializers

[justinchuby]

Fix #1440 by pointing graph input, output and initializers back to the Graph using a tracked list.

Users can now check if a value is a graph input/output/initializer, and find the owning graph of a value with .graph.

[codecov]

❌ 5 Tests Failed:

Tests completed	Failed	Passed	Skipped
15245	5	15240	2521

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

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0829_test_reduce_l2_keep_dims_example

Stack Traces | 0.005s 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_reduce_l2_keep_dims_example'

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_l2_keep_dims_example' (e=No module named 'tests.onnx_backend_test_code.test_reduce_l2_keep_dims_example') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l2_keep_dims_example.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l2_keep_dims_example.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_l2_keep_dims_example(data: FLOAT[3,2,2], axes: INT64[1]) -> (FLOAT[3,2,1]):
E       reduced = opset18.ReduceL2(data, axes, keepdims=1)
E       return reduced

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0012_test_affine_grid_2d

Stack Traces | 0.006s 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_affine_grid_2d'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\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_affine_grid_2d' (e=No module named 'tests.onnx_backend_test_code.test_affine_grid_2d') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_affine_grid_2d.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_affine_grid_2d.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 opset20
E   
E   @script()
E   def bck_test_affine_grid_2d(theta: FLOAT[2,2,3], size: INT64[4]) -> (FLOAT[2,5,6,2]):
E       grid = opset20.AffineGrid(theta, size, align_corners=0)
E       return grid

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1230_test_sce_sum

Stack Traces | 0.006s 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_sce_sum'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\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_sce_sum' (e=No module named 'tests.onnx_backend_test_code.test_sce_sum') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_sce_sum.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_sce_sum.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_sce_sum(x: FLOAT[3,5], y: INT64[3]) -> (FLOAT):
E       z = opset13.SoftmaxCrossEntropyLoss(x, y, reduction='sum')
E       return z

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 updates the IR to record the owning graph for a value when it is used as a graph input/output. The key changes are:

• Introducing tracked container classes (GraphInputs and GraphOutputs) to manage and validate graph inputs/outputs.
• Modifying the Graph initialization to use these tracked containers.
• Adding the owning_graph() method to the Value class to help determine graph ownership.

Reviewed Changes

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

File	Description
onnxscript/ir/_tracked_containers.py	New container classes for graph inputs and outputs with invariance checking and graph assignment logic.
onnxscript/ir/_core.py	Updated Graph initialization to use tracked containers; added owning_graph() method with warning checks.

[Copilot]

Pull Request Overview

This PR addresses issue #1440 by ensuring that graph inputs, outputs, and initializers properly maintain their associations with the owning Graph through a tracked list. Key changes include modifying the constant folding logic to “free” graph outputs by copying and clearing them, updating tests to reflect new graph input behavior, and refactoring initializer handling to avoid incorrect graph ownership.

Reviewed Changes

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

Show a summary per file

File	Description
onnxscript/optimizer/_constant_folding.py	Copies graph outputs into a temporary list and clears the original outputs to allow value movement.
onnxscript/ir/passes/common/constant_manipulation_test.py	Updates unit tests to reflect that certain graph inputs are now removed when initializers are lifted.
onnxscript/ir/passes/common/constant_manipulation.py	Iterates safely over initializers by converting the dict keys to a tuple before modification.
onnxscript/ir/external_data_test.py	Wraps initializer values in new Value objects to reset their graph associations.
onnxscript/ir/_core.py	Updates the Graph API to use the new GraphInput/Output/Initializer containers and adds a .graph property on Value.

Comments suppressed due to low confidence (1)

onnxscript/ir/external_data_test.py:320

• Wrapping the initializer value in a new Value resets its graph association, ensuring proper ownership transfer. Verify that all relevant metadata from the original initializer is preserved in the new Value object.

model.graph.initializers["tensor_same_file"] = ir.Value(