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

Author: justinchuby

onnxscript/ir/_core.py

@@ -31,6 +31,7 @@
     Generic,
     Iterable,
     Iterator,
+    MutableSequence,
     NamedTuple,
     OrderedDict,
     Sequence,
@@ -51,6 +52,7 @@
     _name_authority,
     _protocols,
     _type_casting,
+    _tracked_containers,
 )
 
 if typing.TYPE_CHECKING:
@@ -2116,8 +2118,8 @@ def __init__(
         self.name = name
 
         # Private fields that are not to be accessed by any other classes
-        self._inputs = list(inputs)
-        self._outputs = list(outputs)
+        self._inputs = _tracked_containers.GraphInputs(self, inputs)
+        self._outputs = _tracked_containers.GraphInputs(self, outputs)
         self._initializers = {}
         for initializer in initializers:
             if isinstance(initializer, str):
@@ -2143,11 +2145,11 @@ def __init__(
         self.extend(nodes)
 
     @property
-    def inputs(self) -> list[Value]:
+    def inputs(self) -> MutableSequence[Value]:
         return self._inputs
 
     @property
-    def outputs(self) -> list[Value]:
+    def outputs(self) -> MutableSequence[Value]:
         return self._outputs
 
     @property

onnxscript/ir/_tracked_lists.py renamed to onnxscript/ir/_tracked_containers.py

@@ -1,52 +1,47 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-"""Tracked lists for graph and node IO."""
+"""Tracked containers for graph."""
 
 from __future__ import annotations
 
+__all__ = [
+    "GraphInputs",
+    "GraphOutputs",
+]
+
 import collections
-from typing import TYPE_CHECKING, Iterable, Literal, SupportsIndex
+from typing import TYPE_CHECKING, Iterable, SupportsIndex
+
+import onnxscript
 
 if TYPE_CHECKING:
     from onnxscript.ir import _core
 
 
-class GraphIO(collections.UserList[_core.Value]):
+class _GraphIO(collections.UserList[_core.Value]):
     """The inputs and outputs of a Graph."""
 
-    def __init__(self, graph: _core.Graph, typ: Literal["input", "output"], initlist=None):
+    def __init__(self, graph: _core.Graph, initlist=None):
         super().__init__(initlist)
         self._graph = graph
-        assert typ in {"intput", "output"}
-        self._typ = typ
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        raise NotImplementedError
 
     def _set_graph(self, value: _core.Value) -> None:
         """Set the graph for the value."""
-        if value._graph_input_of is not None and value._graph_input_of is not self._graph:
-            raise ValueError(
-                f"Value '{value}' is already an input of a different graph: {value._graph_input_of!r}"
-            )
-        if value._graph_output_of is not None and value._graph_output_of is not self._graph:
-            raise ValueError(
-                f"Value '{value}' is already an output of a different graph: {value._graph_output_of!r}"
-            )
-
-        if self._typ == "input":
-            value._graph_input_of = self._graph
-        else:
-            value._graph_output_of = self._graph
+        raise NotImplementedError
 
     def _unset_graph(self, value: _core.Value) -> None:
         """Unset the graph for the value."""
-        if self._typ == "input":
-            value._graph_input_of = None
-        else:
-            value._graph_output_of = None
+        raise NotImplementedError
 
     def append(self, item: _core.Value) -> None:
         """Add a new input to the graph."""
         super().append(item)
         self._set_graph(item)
+        self._check_invariance()
 
     def extend(self, other) -> None:
         """Extend the list of inputs or outputs."""
@@ -58,17 +53,20 @@ def insert(self, i: int, item: _core.Value) -> None:
         """Insert an input/output to the graph."""
         super().insert(i, item)
         self._set_graph(item)
+        self._check_invariance()
 
     def pop(self, i: int = -1) -> _core.Value:
         """Remove an input/output from the graph."""
         value = super().pop(i)
         self._unset_graph(value)
+        self._check_invariance()
         return value
 
     def remove(self, item: _core.Value) -> None:
         """Remove an input/output from the graph."""
         super().remove(item)
         self._unset_graph(item)
+        self._check_invariance()
 
     def clear(self) -> None:
         """Clear the list."""
@@ -85,12 +83,76 @@ def __setitem__(self, i, item) -> None:
             for value in item:
                 self._set_graph(value)
             super().__setitem__(i, item)
+            self._check_invariance()
             return
         elif isinstance(item, _core.Value) and isinstance(i, SupportsIndex):
             # Replace a single item
             self._unset_graph(self.data[i])
             self._set_graph(item)
             super().__setitem__(i, item)
+            self._check_invariance()
             return
 
         raise TypeError(f"Invalid types for __setitem__: {type(i)} and {type(item)}")
+
+
+class GraphInputs(_GraphIO):
+    """The inputs of a Graph."""
+
+    def __init__(self, graph: _core.Graph, initlist=None):
+        super().__init__(graph, initlist)
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        if not onnxscript.DEBUG:
+            return
+        for value in self.data:
+            if value._graph_input_of is self._graph:
+                continue
+            raise ValueError(
+                f"Invariance error: Value '{value}' is not an input of the graph: {self._graph!r}"
+            )
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph_input_of is not None and value._graph_input_of is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an input of a different graph: {value._graph_input_of!r}"
+            )
+
+        value._graph_input_of = self._graph
+
+    def _unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        value._graph_input_of = None
+
+
+class GraphOutputs(_GraphIO):
+    """The outputs of a Graph."""
+
+    def __init__(self, graph: _core.Graph, initlist=None):
+        super().__init__(graph, initlist)
+
+    def _check_invariance(self) -> None:
+        """Check the invariance of the graph."""
+        if not onnxscript.DEBUG:
+            return
+        for value in self.data:
+            if value._graph_output_of is self._graph:
+                continue
+            raise ValueError(
+                f"Invariance error: Value '{value}' is not an output of the graph: {self._graph!r}"
+            )
+
+    def _set_graph(self, value: _core.Value) -> None:
+        """Set the graph for the value."""
+        if value._graph_output_of is not None and value._graph_output_of is not self._graph:
+            raise ValueError(
+                f"Value '{value}' is already an output of a different graph: {value._graph_output_of!r}"
+            )
+
+        value._graph_output_of = self._graph
+
+    def _unset_graph(self, value: _core.Value) -> None:
+        """Unset the graph for the value."""
+        value._graph_output_of = None