Rewriter: keep Reshape allowzero valid when materializing the shape

onnxscript/rewriter/rules/common/_materialize_reshape_shape.py

@@ -48,15 +48,24 @@ def check(self, context, data: ir.Value, shape: ir.Value) -> MatchResult:
                 f"Output shape has {sym_count} symbolic dims, cannot materialize."
             )
 
+        # Keep the original allowzero. A materialized 0 dim means "actually zero",
+        # not "copy from input", so it needs allowzero=1, which only exists from
+        # opset 14 and cannot be combined with -1.
+        self._allowzero = context.root.attributes.get("allowzero")
+        if 0 in self._new_dims:
+            if sym_count:
+                return check_result.fail("A 0 dim cannot be combined with -1.")
+            if context.graph_or_function.opset_imports.get("", 0) < 14:
+                return check_result.fail("allowzero requires opset 14 or later.")
+            self._allowzero = 1
+
         return check_result
 
     def rewrite(self, op, data: ir.Value, shape: ir.Value):
         new_shape = op.Constant(
             value=ir.tensor(self._new_dims, dtype=ir.DataType.INT64),
         )
-        # Use allowzero=1 because the materialized dims are concrete values
-        # from shape inference, so a 0 dim means "actually zero" not "copy from input"
-        return op.Reshape(data, new_shape, allowzero=1)
+        return op.Reshape(data, new_shape, allowzero=self._allowzero)
 
 
 materialize_reshape_shape_rule = MaterializeReshapeShape.rule()

onnxscript/rewriter/rules/common/_materialize_reshape_shape_test.py

@@ -136,6 +136,105 @@ def test_allowzero_attribute_preserved(self):
         allowzero = reshape_nodes[0].attributes.get_int("allowzero", 0)
         self.assertEqual(allowzero, 1)
 
+    def test_allowzero_reference_attribute_preserved(self):
+        """An allowzero that refers to a function attribute should be kept as a reference."""
+        model = ir.from_onnx_text(
+            """
+            <ir_version: 10, opset_import: [ "" : 17, "local" : 1]>
+            agraph (float[12] data, float[3, 4] ref) => (float[3, 4] output)
+            {
+                output = local.F<allowzero=0>(data, ref)
+            }
+            <opset_import: [ "" : 17], domain: "local">
+            F<allowzero>(data, ref) => (output)
+            {
+                shape = Shape(ref)
+                output = Reshape<allowzero: int = @allowzero>(data, shape)
+            }
+        """
+        )
+        function = next(iter(model.functions.values()))
+        for node in function:
+            if node.op_type == "Reshape":
+                node.outputs[0].shape = ir.Shape([3, 4])
+                break
+        count = _materialize_reshape_shape.rules.apply_to_model(model)
+        self.assertEqual(count, 1)
+        reshape_nodes = [n for n in function if n.op_type == "Reshape"]
+        self.assertEqual(len(reshape_nodes), 1)
+        allowzero = reshape_nodes[0].attributes.get("allowzero")
+        self.assertIsNotNone(allowzero)
+        self.assertTrue(allowzero.is_ref())
+        self.assertEqual(allowzero.ref_attr_name, "allowzero")
+
+    def test_allowzero_not_added_before_opset_14(self):
+        """Before opset 14 Reshape has no allowzero, so it must not be added."""
+        model_text = """
+            <ir_version: 7, opset_import: [ "" : 13]>
+            agraph (float[12] data, float[3, 4] ref) => (float[3, 4] output)
+            {
+                shape = Shape(ref)
+                output = Reshape(data, shape)
+            }
+        """
+        original = ir.from_onnx_text(model_text)
+        model = ir.from_onnx_text(model_text)
+        for node in model.graph:
+            if node.op_type == "Reshape":
+                node.outputs[0].shape = ir.Shape([3, 4])
+                break
+        count = _materialize_reshape_shape.rules.apply_to_model(model)
+        self.assertEqual(count, 1)
+        reshape_nodes = [n for n in model.graph if n.op_type == "Reshape"]
+        self.assertEqual(len(reshape_nodes), 1)
+        self.assertNotIn("allowzero", reshape_nodes[0].attributes)
+        testing.assert_numerically_equal(
+            original,
+            model,
+            (
+                np.arange(12).astype(np.float32),
+                np.zeros((3, 4), dtype=np.float32),
+            ),
+        )
+
+    def test_zero_dim_not_materialized_before_opset_14(self):
+        """A 0 dim needs allowzero=1, which is not available before opset 14."""
+        model = ir.from_onnx_text(
+            """
+            <ir_version: 7, opset_import: [ "" : 13]>
+            agraph (float[0, 3] data) => (float[0, 3] output)
+            {
+                shape = Shape(data)
+                output = Reshape(data, shape)
+            }
+        """
+        )
+        for node in model.graph:
+            if node.op_type == "Reshape":
+                node.outputs[0].shape = ir.Shape([0, 3])
+                break
+        count = _materialize_reshape_shape.rules.apply_to_model(model)
+        self.assertEqual(count, 0)
+
+    def test_zero_dim_with_symbolic_dim_not_materialized(self):
+        """[0, -1] is invalid with allowzero=1, so the rule should not fire."""
+        model = ir.from_onnx_text(
+            """
+            <ir_version: 7, opset_import: [ "" : 17]>
+            agraph (float[0, N] data) => (float[0, N] output)
+            {
+                shape = Shape(data)
+                output = Reshape(data, shape)
+            }
+        """
+        )
+        for node in model.graph:
+            if node.op_type == "Reshape":
+                node.outputs[0].shape = ir.Shape([0, "N"])
+                break
+        count = _materialize_reshape_shape.rules.apply_to_model(model)
+        self.assertEqual(count, 0)
+
     def test_numerical_correctness_static(self):
         """Verify numerical equivalence for fully static materialization."""
         # Build a model where a dynamic Concat produces the shape for Reshape.