Keep a reducing ScatterND when its static indices cover the whole tensor

onnxscript/rewriter/rules/common/_redundant_scatter_nd.py

@@ -80,9 +80,15 @@ def pattern(self, op, data, indices, updates):
 
     def check(self, context, data, indices, updates, **_):
         """Check if the ScatterND is redundant due to static indices covering entire tensor."""
-        # To validate data can be replaced directly by updates, we need to check the following:
-        # 1. they have the same shape
         result = onnxscript.rewriter.MatchResult()
+        # To validate data can be replaced directly by updates, we need to check the following:
+        # 1. the scatter-update is an assignment, not a reduction combining data and updates
+        reduction = context.root.attributes.get_string("reduction", "none")
+        if reduction != "none":
+            return result.fail(
+                f"The 'reduction' is {reduction!r}, so 'data' still contributes to the result."
+            )
+        # 2. they have the same shape
         if data.shape is None:
             return result.fail("The value 'data' shape is not statically known.", data)
         if updates.shape is None:
@@ -92,7 +98,7 @@ def check(self, context, data, indices, updates, **_):
                 "The shape of 'data' and 'updates' are different.", [data, updates]
             )
 
-        # 2. the indices is referring to the whole data, which is from 0 to data.shape[0]
+        # 3. the indices is referring to the whole data, which is from 0 to data.shape[0]
         if indices.const_value is None:
             return result.fail("The value 'indices' is not statically known.", indices)
         expected_indices = [[i] for i in range(data.shape[0])]

onnxscript/rewriter/rules/common/_redundant_scatter_nd_test.py

@@ -8,6 +8,7 @@
 import onnx.parser
 import onnx_ir as ir
 import onnxruntime
+import parameterized
 from onnx_ir.passes.common import CheckerPass, ShapeInferencePass
 
 import onnxscript.optimizer
@@ -120,6 +121,53 @@ def test_redundant_scatter_nd_static_indices(self):
         for original_output, optimized_output in zip(original_outputs, optimized_outputs):
             np.testing.assert_allclose(original_output, optimized_output, rtol=1e-6, atol=1e-6)
 
+    @parameterized.parameterized.expand(
+        [
+            ("add",),
+            ("mul",),
+            ("max",),
+            ("min",),
+        ]
+    )
+    def test_static_indices_with_reduction_is_not_redundant(self, reduction: str):
+        """A reducing ScatterND combines data with updates, so it is not an Identity."""
+        model_proto = onnx.parser.parse_model(
+            f"""
+            <ir_version: 7, opset_import: [ "" : 18]>
+            agraph (float[8, 16] data, float[8, 16] updates) => (float[8, 16] output)
+            {{
+                output = ScatterND <reduction = "{reduction}"> (data, indices, updates)
+            }}
+        """
+        )
+        indices = np.arange(8).reshape(8, 1).astype(np.int64)
+        model = ir.serde.deserialize_model(model_proto)
+        indices_value = model.graph[0].inputs[1]
+        indices_value.const_value = ir.Tensor(name="indices", value=indices)
+        model.graph.initializers["indices"] = indices_value
+        original_model_proto = ir.serde.serialize_model(model)
+
+        _redundant_scatter_nd.rules.apply_to_model(model)
+        optimized_model_proto = ir.serde.serialize_model(model)
+
+        # Test numerical equivalence
+        inputs = {
+            "data": np.random.rand(8, 16).astype(np.float32),
+            "updates": np.random.rand(8, 16).astype(np.float32),
+        }
+        session = onnxruntime.InferenceSession(
+            original_model_proto.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        original_outputs = session.run(None, inputs)
+        optimized_session = onnxruntime.InferenceSession(
+            optimized_model_proto.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        optimized_outputs = optimized_session.run(None, inputs)
+
+        for original_output, optimized_output in zip(original_outputs, optimized_outputs):
+            np.testing.assert_allclose(original_output, optimized_output, rtol=1e-6, atol=1e-6)
+        self.assertIn("ScatterND", [node.op_type for node in model.graph])
+
 
 if __name__ == "__main__":
     unittest.main()