[IR] Record owning graph for input/output/initializers (#2282)

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`.

Author: justinchuby

onnxscript/ir/_convenience/__init__.py

@@ -321,7 +321,7 @@ def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     Returns:
         A dictionary mapping names to values.
     """
-    values = {}
+    values: dict[str, _core.Value] = {}
     values.update(graph.initializers)
     # The names of the values can be None or "", which we need to exclude
     for input in graph.inputs:

onnxscript/ir/_core.py

@@ -31,6 +31,8 @@
     Generic,
     Iterable,
     Iterator,
+    MutableMapping,
+    MutableSequence,
     NamedTuple,
     OrderedDict,
     Sequence,
@@ -46,6 +48,7 @@
 from onnxscript.ir import (
     _display,
     _enums,
+    _graph_containers,
     _linked_list,
     _metadata,
     _name_authority,
@@ -1746,18 +1749,19 @@ class Value(_protocols.ValueProtocol, _display.PrettyPrintable):
 
     To find all the nodes that use this value as an input, call :meth:`uses`.
 
-    To check if the value is an output of a graph, call :meth:`is_graph_output`.
+    To check if the value is an is an input, output or initializer of a graph,
+    use :meth:`is_graph_input`, :meth:`is_graph_output` or :meth:`is_initializer`.
 
-    Attributes:
-        name: The name of the value. A value is always named when it is part of a graph.
-        shape: The shape of the value.
-        type: The type of the value.
-        metadata_props: Metadata.
+    Use :meth:`graph` to get the graph that owns the value.
     """
 
     __slots__ = (
         "_const_value",
+        "_graph",
         "_index",
+        "_is_graph_input",
+        "_is_graph_output",
+        "_is_initializer",
         "_metadata",
         "_metadata_props",
         "_name",
@@ -1808,6 +1812,14 @@ def __init__(
         self._uses: dict[Usage, None] = {}
         self.doc_string = doc_string
 
+        # The graph this value belongs to. It is set *only* when the value is added as
+        # a graph input, output or initializer.
+        # The four properties can only be set by the Graph class (_GraphIO and GraphInitializers).
+        self._graph: Graph | None = None
+        self._is_graph_input: bool = False
+        self._is_graph_output: bool = False
+        self._is_initializer: bool = False
+
     def __repr__(self) -> str:
         value_name = self.name if self.name else "anonymous:" + str(id(self))
         type_text = f", type={self.type!r}" if self.type is not None else ""
@@ -1846,11 +1858,35 @@ def _constant_tensor_part(self) -> str:
                 return f"{{{self.const_value.__class__.__name__}(...)}}"
         return ""
 
+    @property
+    def graph(self) -> Graph | None:
+        """Return the graph that defines this value.
+
+        When the value is an input/output/initializer of a graph, the owning graph
+        is that graph. When the value is an output of a node, the owning graph is the
+        graph that the node belongs to. When the value is not owned by any graph,
+        it returns ``None``.
+        """
+        if self._graph is not None:
+            return self._graph
+        if self._producer is not None:
+            return self._producer.graph
+        return None
+
+    def _owned_by_graph(self) -> bool:
+        """Return True if the value is owned by a graph."""
+        result = self._is_graph_input or self._is_graph_output or self._is_initializer
+        if result:
+            assert self._graph is not None
+        return result
+
     def producer(self) -> Node | None:
         """The node that produces this value.
 
         When producer is ``None``, the value does not belong to a node, and is
-        typically a graph input or an initializer.
+        typically a graph input or an initializer. You can use :meth:`graph``
+        to find the graph that owns this value. Use :meth:`is_graph_input`, :meth:`is_graph_output`
+        or :meth:`is_initializer` to check if the value is an input, output or initializer of a graph.
         """
         return self._producer
 
@@ -1986,15 +2022,17 @@ def metadata_props(self) -> dict[str, str]:
             self._metadata_props = {}
         return self._metadata_props
 
+    def is_graph_input(self) -> bool:
+        """Whether the value is an input of a graph."""
+        return self._is_graph_input
+
     def is_graph_output(self) -> bool:
         """Whether the value is an output of a graph."""
-        if (producer := self.producer()) is None:
-            return False
-        if (graph := producer.graph) is None:
-            return False
-        # Cannot use `in` because __eq__ may be defined by subclasses, even though
-        # it is not recommended
-        return any(output is self for output in graph.outputs)
+        return self._is_graph_output
+
+    def is_initializer(self) -> bool:
+        """Whether the value is an initializer of a graph."""
+        return self._is_initializer
 
 
 def Input(
@@ -2104,9 +2142,9 @@ 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._initializers = {}
+        self._inputs = _graph_containers.GraphInputs(self, inputs)
+        self._outputs = _graph_containers.GraphOutputs(self, outputs)
+        self._initializers = _graph_containers.GraphInitializers(self)
         for initializer in initializers:
             if isinstance(initializer, str):
                 raise TypeError(
@@ -2131,15 +2169,15 @@ 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
-    def initializers(self) -> dict[str, Value]:
+    def initializers(self) -> MutableMapping[str, Value]:
         return self._initializers
 
     def register_initializer(self, value: Value) -> None:
@@ -2159,15 +2197,15 @@ def register_initializer(self, value: Value) -> None:
             ValueError: If the initializer is produced by a node.
             ValueError: If the value does not have its ``.const_value`` set.
         """
+        if not value.name:
+            raise ValueError(f"Initializer must have a name: {value!r}")
         if value.name in self._initializers:
             if self._initializers[value.name] is not value:
                 raise ValueError(
                     f"Initializer '{value.name}' is already registered, but"
                     " it is not the same object: existing={self._initializers[value.name]!r},"
                     f" new={value!r}"
                 )
-        if not value.name:
-            raise ValueError(f"Initializer must have a name: {value!r}")
         if value.producer() is not None:
             raise ValueError(
                 f"Value '{value!r}' is produced by a node and cannot be an initializer."
@@ -2858,11 +2896,11 @@ def overload(self, value: str) -> None:
         self._overload = value
 
     @property
-    def inputs(self) -> list[Value]:
+    def inputs(self) -> MutableSequence[Value]:
         return self._graph.inputs
 
     @property
-    def outputs(self) -> list[Value]:
+    def outputs(self) -> MutableSequence[Value]:
         return self._graph.outputs
 
     @property