Fix metadata props handling for values (#198)

1. Add metadata_props to value initializers
2. Fix serde to persist the field correctly

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/onnx_ir/_convenience/_constructors.py

@@ -224,6 +224,7 @@ def val(
     *,
     type: ir.TypeProtocol | None = None,
     const_value: ir.TensorProtocol | None = None,
+    metadata_props: dict[str, str] | None = None,
 ) -> ir.Value:
     """Create a :class:`~onnx_ir.Value` with the given name and type.
 
@@ -253,6 +254,7 @@ def val(
         type: The type of the value. Only one of dtype and type can be specified.
         const_value: The constant tensor that initializes the value. Supply this argument
             when you want to create an initializer. The type and shape can be obtained from the tensor.
+        metadata_props: The metadata properties that will be serialized to the ONNX proto.
 
     Returns:
         A Value object.
@@ -279,10 +281,11 @@ def val(
             type=const_tensor_type,
             shape=_core.Shape(const_value.shape),  # type: ignore
             const_value=const_value,
+            metadata_props=metadata_props,
         )
 
     if type is None and dtype is not None:
         type = _core.TensorType(dtype)
     if shape is not None and not isinstance(shape, _core.Shape):
         shape = _core.Shape(shape)
-    return _core.Value(name=name, type=type, shape=shape)
+    return _core.Value(name=name, type=type, shape=shape, metadata_props=metadata_props)

src/onnx_ir/_core.py

@@ -165,6 +165,11 @@ def nbytes(self) -> int:
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the tensor.
+
+        The metadata properties are used to store additional information about the tensor.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -2022,6 +2027,7 @@ def __init__(
         type: _protocols.TypeProtocol | None = None,
         doc_string: str | None = None,
         const_value: _protocols.TensorProtocol | None = None,
+        metadata_props: dict[str, str] | None = None,
     ) -> None:
         """Initialize a value.
 
@@ -2034,11 +2040,12 @@ def __init__(
             type: The type of the value.
             doc_string: The documentation string.
             const_value: The constant tensor if the value is constant.
+            metadata_props: Metadata that will be serialized to the ONNX file.
         """
         self._producer: Node | None = producer
         self._index: int | None = index
         self._metadata: _metadata.MetadataStore | None = None
-        self._metadata_props: dict[str, str] | None = None
+        self._metadata_props: dict[str, str] | None = metadata_props
 
         self._name: str | None = name
         self._shape: Shape | None = shape
@@ -2226,9 +2233,16 @@ def shape(self, value: Shape | None) -> None:
     def const_value(
         self,
     ) -> _protocols.TensorProtocol | None:
-        """A concrete value.
+        """The backing constant tensor for the value.
 
-        The value can be backed by different raw data types, such as numpy arrays.
+        If the ``Value`` has a ``const_value`` and is part of a graph initializers
+        dictionary, the value is an initialized value. Its ``const_value``
+        will appear as an ``initializer`` in the GraphProto when serialized.
+
+        If the ``Value`` is not part of a graph initializers dictionary, the ``const_value``
+        field will be ignored during serialization.
+
+        ``const_value`` can be backed by different raw data types, such as numpy arrays.
         The only guarantee is that it conforms TensorProtocol.
         """
         return self._const_value
@@ -2258,6 +2272,11 @@ def meta(self) -> _metadata.MetadataStore:
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the value.
+
+        The metadata properties are used to store additional information about the value.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -2805,6 +2824,11 @@ def meta(self) -> _metadata.MetadataStore:
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the graph.
+
+        The metadata properties are used to store additional information about the graph.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -3057,6 +3081,11 @@ def meta(self) -> _metadata.MetadataStore:
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the model.
+
+        The metadata properties are used to store additional information about the model.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         if self._metadata_props is None:
             self._metadata_props = {}
         return self._metadata_props
@@ -3250,6 +3279,11 @@ def meta(self) -> _metadata.MetadataStore:
 
     @property
     def metadata_props(self) -> dict[str, str]:
+        """The metadata properties of the function.
+
+        The metadata properties are used to store additional information about the function.
+        Unlike ``meta``, this property is serialized to the ONNX proto.
+        """
         return self._graph.metadata_props
 
     def all_nodes(self) -> Iterator[Node]:

src/onnx_ir/serde.py

@@ -709,8 +709,7 @@ def _deserialize_graph(
         annotation.tensor_name: annotation for annotation in proto.quantization_annotation
     }
 
-    # Create values for initializers and inputs
-    initializer_tensors = [deserialize_tensor(tensor) for tensor in proto.initializer]
+    # Create values for inputs
     inputs = [_core.Value(name=info.name) for info in proto.input]
     for info, value in zip(proto.input, inputs):
         deserialize_value_info_proto(info, value)
@@ -725,6 +724,11 @@ def _deserialize_graph(
     # Enter the graph scope by pushing the values for this scope to the stack
     scoped_values.append(values)
 
+    # Build the value info dictionary to allow for quick lookup for this graph scope
+    value_info = {info.name: info for info in proto.value_info}
+
+    # Create values for initializers
+    initializer_tensors = [deserialize_tensor(tensor) for tensor in proto.initializer]
     initializer_values = []
     for i, tensor in enumerate(initializer_tensors):
         initializer_name = tensor.name
@@ -750,16 +754,15 @@ def _deserialize_graph(
                 shape=tensor.shape,  # type: ignore[arg-type]
                 const_value=tensor,
             )
+            if initializer_name in value_info:
+                deserialize_value_info_proto(value_info[initializer_name], initializer_value)
             if initializer_value.name in quantization_annotations:
                 _deserialize_quantization_annotation(
                     quantization_annotations[initializer_value.name], initializer_value
                 )
             values[initializer_name] = initializer_value
         initializer_values.append(initializer_value)
 
-    # Build the value info dictionary to allow for quick lookup for this graph scope
-    value_info = {info.name: info for info in proto.value_info}
-
     # Declare values for all node outputs from this graph scope. This is necessary
     # to handle the case where a node in a subgraph uses a value that is declared out
     # of order in the outer graph. Declaring the values first allows us to find the
@@ -1390,7 +1393,12 @@ def _should_create_value_info_for_value(value: _protocols.ValueProtocol) -> bool
         True if value info should be created for the value.
     """
     # No need to serialize value info if it is not set
-    if value.shape is None and value.type is None:
+    if (
+        value.shape is None
+        and value.type is None
+        and not value.metadata_props
+        and not value.doc_string
+    ):
         return False
     if not value.name:
         logger.debug("Did not serialize '%s' because its name is empty", value)

src/onnx_ir/serde_test.py

@@ -531,6 +531,45 @@ def test_deserialize_builds_correct_value_connections_for_subgraphs_that_referen
             [n.name for n in deserialized_model.graph], ["b_producer", "node_with_subgraph"]
         )
 
+    def test_value_metadata_props_are_preserved(self):
+        value = ir.val(
+            "test_initializer",
+            dtype=ir.DataType.FLOAT,
+            shape=(2,),
+            const_value=ir.tensor([1.0, 2.0], name="test_initializer"),
+            metadata_props={"key": "value"},
+        )
+        input = ir.val(
+            "test_input", dtype=ir.DataType.FLOAT, shape=(2,), metadata_props={"key": "input"}
+        )
+        node = ir.node("Identity", inputs=[input])
+        node.outputs[0].metadata_props["key"] = "intermediate"
+        output = ir.val(
+            "test_output",
+            dtype=ir.DataType.FLOAT,
+            shape=(2,),
+            metadata_props={"key": "output"},
+        )
+        node2 = ir.node("Identity", inputs=node.outputs, outputs=[output])
+        graph = ir.Graph(
+            inputs=[input],
+            outputs=[output],
+            nodes=[node, node2],
+            initializers=[value],
+            name="test_graph",
+        )
+        graph_proto = serde.serialize_graph(graph)
+        deserialized_graph = serde.deserialize_graph(graph_proto)
+
+        self.assertEqual(deserialized_graph.inputs[0].metadata_props, {"key": "input"})
+        self.assertEqual(deserialized_graph.outputs[0].metadata_props, {"key": "output"})
+        intermediate_value = deserialized_graph.node(0).outputs[0]
+        self.assertEqual(intermediate_value.metadata_props, {"key": "intermediate"})
+
+        self.assertIn("test_initializer", deserialized_graph.initializers)
+        deserialized_value = deserialized_graph.initializers["test_initializer"]
+        self.assertEqual(deserialized_value.metadata_props, {"key": "value"})
+
 
 class SerializationTest(unittest.TestCase):
     @parameterized.parameterized.expand(