Ensure int/float attributes are stored as Python int/float (#178)

Ensure int/float attributes are stored as Python int/float by doing a
quick type conversion at Attr initialization. Even though this makes
Attr initialization slightly more costly, it allows us to maintain an
invariance where an INT or FLOAT value can only be a plain python
number, making downstream usage and comparison much easier.

---------

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

Author: justinchuby

src/onnx_ir/_convenience/__init__.py

@@ -226,7 +226,7 @@ def convert_attributes(
         ...     "type_protos": [ir.TensorType(ir.DataType.FLOAT), ir.TensorType(ir.DataType.FLOAT)],
         ... }
         >>> convert_attributes(attrs)
-        [Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', GRAPH, Graph(
+        [Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, (1, 2, 3)), Attr('floats', FLOATS, (1.0, 2.0, 3.0)), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', GRAPH, Graph(
             name='graph0',
             inputs=(
         <BLANKLINE>

src/onnx_ir/_core.py

@@ -3397,6 +3397,24 @@ def __init__(
         *,
         doc_string: str | None = None,
     ) -> None:
+        # Quick checks to ensure that INT and FLOAT attributes are stored as int and float,
+        # not np.int32, np.float32, bool, etc.
+        # This also allows errors to be raised at the time of construction instead of later
+        # during serialization.
+        # TODO(justinchuby): Use case matching when we drop support for Python 3.9
+        if value is None:
+            # Value can be None for reference attributes or when it is used as a
+            # placeholder for schemas
+            pass
+        elif type == _enums.AttributeType.INT:
+            value = int(value)
+        elif type == _enums.AttributeType.FLOAT:
+            value = float(value)
+        elif type == _enums.AttributeType.INTS:
+            value = tuple(int(v) for v in value)
+        elif type == _enums.AttributeType.FLOATS:
+            value = tuple(float(v) for v in value)
+
         self._name = name
         self._type = type
         self._value = value
@@ -3472,17 +3490,17 @@ def as_float(self) -> float:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type FLOAT. Actual type: {self.type}"
             )
-        # Do not use isinstance check because it may prevent np.float32 etc. from being used
-        return float(self.value)
+        # value is guaranteed to be a float in the constructor
+        return self.value
 
     def as_int(self) -> int:
         """Get the attribute value as an int."""
         if self.type != _enums.AttributeType.INT:
             raise TypeError(
                 f"Attribute '{self.name}' is not of type INT. Actual type: {self.type}"
             )
-        # Do not use isinstance check because it may prevent np.int32 etc. from being used
-        return int(self.value)
+        # value is guaranteed to be an int in the constructor
+        return self.value
 
     def as_string(self) -> str:
         """Get the attribute value as a string."""
@@ -3522,9 +3540,8 @@ def as_floats(self) -> Sequence[float]:
             )
         if not isinstance(self.value, Sequence):
             raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
-        # Do not use isinstance check on elements because it may prevent np.int32 etc. from being used
-        # Create a copy of the list to prevent mutation
-        return [float(v) for v in self.value]
+        # value is guaranteed to be a sequence of float in the constructor
+        return self.value
 
     def as_ints(self) -> Sequence[int]:
         """Get the attribute value as a sequence of ints."""
@@ -3534,9 +3551,8 @@ def as_ints(self) -> Sequence[int]:
             )
         if not isinstance(self.value, Sequence):
             raise TypeError(f"Value of attribute '{self!r}' is not a Sequence.")
-        # Do not use isinstance check on elements because it may prevent np.int32 etc. from being used
-        # Create a copy of the list to prevent mutation
-        return list(self.value)
+        # value is guaranteed to be a sequence of int in the constructor
+        return self.value
 
     def as_strings(self) -> Sequence[str]:
         """Get the attribute value as a sequence of strings."""
@@ -3605,7 +3621,7 @@ def RefAttr(
     return Attr(name, type, None, ref_attr_name=ref_attr_name, doc_string=doc_string)
 
 
-def AttrFloat32(name: str, value: float, doc_string: str | None = None) -> Attr:
+def AttrFloat32(name: str, value: float | np.floating, doc_string: str | None = None) -> Attr:
     """Create a float attribute."""
     # NOTE: The function name is capitalized to maintain API backward compatibility.
     return Attr(
@@ -3616,7 +3632,7 @@ def AttrFloat32(name: str, value: float, doc_string: str | None = None) -> Attr:
     )
 
 
-def AttrInt64(name: str, value: int, doc_string: str | None = None) -> Attr:
+def AttrInt64(name: str, value: int | np.integer, doc_string: str | None = None) -> Attr:
     """Create an int attribute."""
     # NOTE: The function name is capitalized to maintain API backward compatibility.
     return Attr(

src/onnx_ir/_core_test.py

@@ -950,7 +950,7 @@ def test_attributes_get_ints(self):
             inputs=(),
             attributes=[_core.AttrInt64s("test_attr", [1, 2, 3])],
         )
-        self.assertEqual(node.attributes.get_ints("test_attr"), [1, 2, 3])
+        self.assertEqual(node.attributes.get_ints("test_attr"), (1, 2, 3))
         self.assertIsNone(node.attributes.get_ints("non_existent_attr"))
         self.assertEqual(node.attributes.get_ints("non_existent_attr", [42]), [42])
 
@@ -961,7 +961,7 @@ def test_attributes_get_floats(self):
             inputs=(),
             attributes=[_core.AttrFloat32s("test_attr", [1.0, 2.0, 3.0])],
         )
-        self.assertEqual(node.attributes.get_floats("test_attr"), [1.0, 2.0, 3.0])
+        self.assertEqual(node.attributes.get_floats("test_attr"), (1.0, 2.0, 3.0))
         self.assertIsNone(node.attributes.get_floats("non_existent_attr"))
         self.assertEqual(node.attributes.get_floats("non_existent_attr", [42.0]), [42.0])
 
@@ -1971,11 +1971,11 @@ def test_as_graph(self):
 
     def test_as_floats(self):
         attr = _core.Attr("test", ir.AttributeType.FLOATS, [42.0])
-        self.assertEqual(attr.as_floats(), [42.0])
+        self.assertEqual(tuple(attr.as_floats()), (42.0,))
 
     def test_as_ints(self):
         attr = _core.Attr("test", ir.AttributeType.INTS, [42])
-        self.assertEqual(attr.as_ints(), [42])
+        self.assertEqual(tuple(attr.as_ints()), (42,))
 
     def test_as_strings(self):
         attr = _core.Attr("test", ir.AttributeType.STRINGS, ["test string", ""])

src/onnx_ir/serde.py

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

src/onnx_ir/serde_test.py

@@ -539,8 +539,9 @@ class SerializationTest(unittest.TestCase):
             ("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]),
+            ("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)),
+            ("bools", ir.AttributeType.INTS, [True, False], (1, 0)),
             ("string", ir.AttributeType.STRING, "test_string", "test_string"),
         ]
     )