[pass] Preserve value metadata in LiftConstantsToInitializersPass (#204)

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

Author: justinchuby

src/onnx_ir/passes/common/constant_manipulation.py

@@ -69,7 +69,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                 shape=tensor.shape,  # type: ignore[arg-type]
                 type=ir.TensorType(tensor.dtype),
                 const_value=tensor,
+                # Preserve metadata from Constant value into the onnx model
+                metadata_props=node.outputs[0].metadata_props.copy(),
             )
+            # Preserve value meta from the Constant output for intermediate analysis
+            initializer.meta.update(node.outputs[0].meta)
+
             assert node.graph is not None
             node.graph.register_initializer(initializer)
             # Replace the constant node with the initializer

src/onnx_ir/passes/common/constant_manipulation_test.py

@@ -36,6 +36,8 @@ def test_pass_with_lifting_float_and_int_constants_to_initializers(
         const_node = ir.node(
             "Constant", inputs=[], attributes={"value": constant_tensor}, num_outputs=1
         )
+        const_node.outputs[0].meta["meta_key"] = "meta_val"
+        const_node.outputs[0].metadata_props["metadata_key"] = "metadata_val"
         add_node = ir.node("Add", inputs=[inputs[0], const_node.outputs[0]])
         mul_node = ir.node("Mul", inputs=[add_node.outputs[0], inputs[1]])
 
@@ -72,6 +74,15 @@ def test_pass_with_lifting_float_and_int_constants_to_initializers(
         self.assertEqual(
             len([node for node in result.model.graph if node.op_type == "Constant"]), 0
         )
+        # Metadata is preserved
+        self.assertEqual(
+            result.model.graph.initializers["val_0"].meta,
+            {"meta_key": "meta_val"},
+        )
+        self.assertEqual(
+            result.model.graph.initializers["val_0"].metadata_props,
+            {"metadata_key": "metadata_val"},
+        )
 
     @parameterized.parameterized.expand(
         [