diff --git a/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py b/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py index 8c63d3e777..7cae5cd2c5 100644 --- a/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py +++ b/onnxscript/rewriter/rules/common/_cast_constant_of_shape.py @@ -4,6 +4,7 @@ import logging +import numpy as np import onnx_ir as ir from onnxscript.rewriter._rewrite_rule import RewriteRule, RewriteRuleSet @@ -17,9 +18,13 @@ def cast_constant_of_shape(op, shape, scalar, dtype): def fused_cast_constant_of_shape(op, shape: ir.Value, scalar: ir.Attr, dtype: ir.Attr, **_): - # Cast scalar (a TensorProto attribute) to the specified dtype - scalar_value = scalar.value.numpy().item() - cast_value = ir.tensor([scalar_value], dtype=ir.DataType(dtype.as_int())) + # Cast scalar (a TensorProto attribute) to the specified dtype. Convert the array, not a + # Python scalar, so that values outside the target range wrap like the Cast op instead of + # raising OverflowError. + to_dtype = ir.DataType(dtype.as_int()) + with np.errstate(invalid="ignore"): + cast_array = scalar.value.numpy().astype(to_dtype.numpy()).reshape(1) + cast_value = ir.tensor(cast_array, dtype=to_dtype) return op.ConstantOfShape(shape, value=cast_value) diff --git a/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py b/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py index 9001f95806..bccba0a4ad 100644 --- a/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py +++ b/onnxscript/rewriter/rules/common/_cast_constant_of_shape_test.py @@ -2,9 +2,11 @@ # Licensed under the MIT License. import unittest +import numpy as np import onnx.checker import onnx.parser import onnx_ir as ir +import parameterized from onnxscript.rewriter.rules.common import _cast_constant_of_shape @@ -49,6 +51,34 @@ def test_cast_after_constant_of_shape_without_value_is_fused(self): output_model_proto = ir.serde.serialize_model(model) onnx.checker.check_model(output_model_proto, True) + @parameterized.parameterized.expand( + [ + ("int_above_uint8_range", "int32", "{300}", 2, np.array([300], np.int32)), + ("negative_int_to_uint8", "int32", "{-1}", 2, np.array([-1], np.int32)), + ("int64_to_int8", "int64", "{200}", 3, np.array([200], np.int64)), + ] + ) + def test_fused_value_matches_cast_of_out_of_range_constant( + self, _, value_type, value, to, value_array + ): + model_proto = onnx.parser.parse_model( + f""" + + agraph (int64[2] input_x) => (output) + {{ + constant = ConstantOfShape (input_x) + output = Cast (constant) + }} + """ + ) + model = ir.serde.deserialize_model(model_proto) + count = _cast_constant_of_shape.rules.apply_to_model(model) + self.assertEqual(count, 1) + fused_value = model.graph[0].attributes["value"].value + self.assertEqual(fused_value.dtype, to) + expected = value_array.astype(ir.DataType(to).numpy()) + np.testing.assert_array_equal(fused_value.numpy(), expected) + if __name__ == "__main__": unittest.main()