[serde] Cast attribute values to int or float (#175)

Cast attribute values to int or float based on the attribute type when
serializing. Otherwise protobuf will complain

> Field onnx.AttributeProto.i: Expected an int, got a boolean. This will
be rejected in 7.34.0, please fix it before that


This should fix pytorch/pytorch#161941.

---------

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

Author: justinchuby

src/onnx_ir/serde.py

@@ -1784,11 +1784,12 @@ def _fill_in_value_for_attribute(
 ) -> None:
     if type_ == _enums.AttributeType.INT:
         # value: int
-        attribute_proto.i = value
+        # Cast bool to int, for example
+        attribute_proto.i = int(value)
         attribute_proto.type = onnx.AttributeProto.INT
     elif type_ == _enums.AttributeType.FLOAT:
         # value: float
-        attribute_proto.f = value
+        attribute_proto.f = float(value)
         attribute_proto.type = onnx.AttributeProto.FLOAT
     elif type_ == _enums.AttributeType.STRING:
         # value: str

src/onnx_ir/serde_test.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import itertools
 import unittest
+import warnings
 
 import google.protobuf.text_format
 import ml_dtypes
@@ -531,6 +532,33 @@ def test_deserialize_builds_correct_value_connections_for_subgraphs_that_referen
         )
 
 
+class SerializationTest(unittest.TestCase):
+    @parameterized.parameterized.expand(
+        [
+            ("float", ir.AttributeType.FLOAT, 1.5, 1.5),
+            ("int_as_float", ir.AttributeType.FLOAT, 1, 1.0),
+            ("int", ir.AttributeType.INT, 42, 42),
+            ("bool", ir.AttributeType.INT, True, 1),
+            ("ints", ir.AttributeType.INTS, [1, 2, 3], [1, 2, 3]),
+            ("floats", ir.AttributeType.FLOATS, [1.0, 2.0, 3.0], [1.0, 2.0, 3.0]),
+            ("string", ir.AttributeType.STRING, "test_string", "test_string"),
+        ]
+    )
+    def test_serialize_attribute(self, _: str, typ: ir.AttributeType, value, expected):
+        attr = ir.Attr("test_attr", typ, value)
+        with warnings.catch_warnings(record=True) as w:
+            # Ensure all warnings are caught, not just the default ones
+            warnings.simplefilter("always")
+            attr_proto = serde.serialize_attribute(attr)
+            self.assertEqual(
+                len(w), 0, f"Unexpected warnings: {[str(warn.message) for warn in w]}"
+            )
+        deserialized_attr = serde.deserialize_attribute(attr_proto)
+        self.assertEqual(deserialized_attr.name, attr.name)
+        self.assertEqual(deserialized_attr.type, attr.type)
+        self.assertEqual(deserialized_attr.value, expected)
+
+
 class QuantizationAnnotationTest(unittest.TestCase):
     """Test that quantization annotations are correctly serialized and deserialized."""