Update from_onnx_text method (#80)

- Change the initializers parameter to take a list of tensors.
- Add the missing tests and expose the `to_onnx_text` method in the root
namespace.
- Update `create_value_mapping` to include values from subgraphs.

---------

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/__init__.py

@@ -73,6 +73,7 @@
     "from_proto",
     "from_onnx_text",
     "to_proto",
+    "to_onnx_text",
     # Convenience constructors
     "tensor",
     "node",
@@ -149,7 +150,7 @@
     TypeProtocol,
     ValueProtocol,
 )
-from onnx_ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_proto
+from onnx_ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_onnx_text, to_proto
 
 DEBUG = False
 

src/onnx_ir/_convenience/__init__.py

@@ -21,7 +21,7 @@
 
 import onnx
 
-from onnx_ir import _core, _enums, _protocols, serde
+from onnx_ir import _core, _enums, _protocols, serde, traversal
 
 SupportedAttrTypes = Union[
     str,
@@ -313,7 +313,9 @@ def replace_all_uses_with(
 def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     """Return a dictionary mapping names to values in the graph.
 
-    The mapping does not include values from subgraphs.
+    The mapping includes values from subgraphs. Duplicated names are omitted,
+    and the first value with that name is returned. Values with empty names
+    are excluded from the mapping.
 
     Args:
         graph: The graph to extract the mapping from.
@@ -327,11 +329,23 @@ def create_value_mapping(graph: _core.Graph) -> dict[str, _core.Value]:
     for input in graph.inputs:
         if not input.name:
             continue
+        if input.name in values:
+            continue
         values[input.name] = input
-    for node in graph:
+    for node in traversal.RecursiveGraphIterator(graph):
+        for value in node.inputs:
+            if not value:
+                continue
+            if not value.name:
+                continue
+            if value.name in values:
+                continue
+            values[value.name] = value
         for value in node.outputs:
             if not value.name:
                 continue
+            if value.name in values:
+                continue
             values[value.name] = value
     return values
 

src/onnx_ir/serde.py

@@ -63,7 +63,7 @@
 import collections
 import logging
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, Callable
 
 import numpy as np
@@ -194,30 +194,36 @@ def from_proto(proto: object) -> object:
 def from_onnx_text(
     model_text: str,
     /,
-    with_initializers: Mapping[str, _protocols.TensorProtocol] | None = None,
+    initializers: Iterable[_protocols.TensorProtocol] | None = None,
 ) -> _core.Model:
     """Convert the ONNX textual representation to an IR model.
 
     Read more about the textual representation at: https://onnx.ai/onnx/repo-docs/Syntax.html
 
     Args:
         model_text: The ONNX textual representation of the model.
-        with_initializers: A mapping of initializer names to tensors. If provided, these tensors
+        initializers: Tensors to be added as initializers. If provided, these tensors
             will be added to the model as initializers. If a name does not exist in the model,
             a ValueError will be raised.
 
     Returns:
         The IR model corresponding to the ONNX textual representation.
 
     Raises:
-        ValueError: If a name in `with_initializers` does not exist in the model.
+        ValueError: If a tensor name in `initializers` does not match any value in the model.
     """
     proto = onnx.parser.parse_model(model_text)
     model = deserialize_model(proto)
     values = _convenience.create_value_mapping(model.graph)
-    if with_initializers:
+    if initializers:
         # Add initializers to the model
-        for name, tensor in with_initializers.items():
+        for tensor in initializers:
+            name = tensor.name
+            if not name:
+                raise ValueError(
+                    "Initializer tensor must have a name. "
+                    f"Please provide a name for the initializer: {tensor}"
+                )
             if name not in values:
                 raise ValueError(f"Value '{name}' does not exist in model.")
             initializer = values[name]

src/onnx_ir/serde_test.py

@@ -57,7 +57,7 @@ def test_to_proto(self, _: str, ir_object):
     def test_from_to_onnx_text(self):
         model_text = """\
 <
-   ir_version: 7,
+   ir_version: 10,
    opset_import: ["" : 17]
 >
 agraph (float[1,4,512,512] input_x, float[1,4,512,64] input_y) => (float[4,512,512] reshape_x) {
@@ -67,12 +67,79 @@ def test_from_to_onnx_text(self):
         self.maxDiff = None
         model = serde.from_onnx_text(model_text)
         self.assertIsInstance(model, ir.Model)
-        self.assertEqual(model.ir_version, 7)
+        self.assertEqual(model.ir_version, 10)
         self.assertEqual(len(model.graph.inputs), 2)
         self.assertEqual(len(model.graph.outputs), 1)
         onnx_text_roundtrip = serde.to_onnx_text(model)
         self.assertEqual(model_text.strip(), onnx_text_roundtrip.strip())
 
+    def test_from_to_onnx_text_with_initializers(self):
+        model_text = """\
+<
+   ir_version: 10,
+   opset_import: ["" : 17]
+>
+agraph (float[1] input_x, float[2] input_y) => (float[2] result) {
+   [node_1] add = Add (input_x, input_y)
+   [node_2] result = Add (add, initializer_z)
+}"""
+        self.maxDiff = None
+        array = np.array([1.0, 2.0], dtype=np.float32)
+        init_array = np.array([3.0, 4.0], dtype=np.float32)
+        model = serde.from_onnx_text(
+            model_text,
+            initializers=[
+                ir.tensor(init_array, name="initializer_z"),
+                ir.tensor(array, name="input_y"),
+            ],
+        )
+        np.testing.assert_array_equal(model.graph.inputs[1].const_value.numpy(), array)
+        np.testing.assert_array_equal(
+            model.graph.initializers["initializer_z"].const_value.numpy(), init_array
+        )
+        expected_text = """\
+<
+   ir_version: 10,
+   opset_import: ["" : 17]
+>
+agraph (float[1] input_x, float[2] input_y) => (float[2] result)
+   <float[2] initializer_z =  {3,4}, float[2] input_y =  {1,2}>
+{
+   [node_1] add = Add (input_x, input_y)
+   [node_2] result = Add (add, initializer_z)
+}"""
+        onnx_text_roundtrip = serde.to_onnx_text(model)
+        stripped_lines = [line.rstrip() for line in onnx_text_roundtrip.splitlines()]
+        result = "\n".join(stripped_lines)
+        self.assertEqual(result, expected_text)
+
+    def test_to_onnx_text_excluding_initializers(self):
+        model_text = """\
+<
+   ir_version: 10,
+   opset_import: ["" : 17]
+>
+agraph (float[1] input_x, float[2] input_y) => (float[2] result) {
+   [node_name] result = Add (input_x, input_y)
+}"""
+        self.maxDiff = None
+        array = np.array([1.0, 2.0], dtype=np.float32)
+        model = serde.from_onnx_text(
+            model_text, initializers=[ir.tensor(array, name="input_y")]
+        )
+        onnx_text_without_initializers = serde.to_onnx_text(model, exclude_initializers=True)
+        expected_text_without_initializers = """\
+<
+   ir_version: 10,
+   opset_import: ["" : 17]
+>
+agraph (float[1] input_x, float[2] input_y) => (float[2] result) {
+   [node_name] result = Add (input_x, input_y)
+}"""
+        self.assertEqual(
+            onnx_text_without_initializers.strip(), expected_text_without_initializers
+        )
+
 
 class TensorProtoTensorTest(unittest.TestCase):
     @parameterized.parameterized.expand(