Fix include_self for scatter_reduce

onnxscript/function_libs/torch_lib/ops/core.py

@@ -14,6 +14,9 @@
 import math
 from typing import Any, Optional, Sequence, Tuple, Union
 
+import numpy as np
+import ml_dtypes
+
 from onnxscript import (
     BFLOAT16,
     BOOL,
@@ -7589,7 +7592,67 @@ def aten_scatter_reduce(
         self = op.Reshape(self, neg_1)
         index = op.Reshape(index, neg_1)
         src = op.Reshape(src, neg_1)
+
+    if not include_self:
+        # onnx standard always assume the value from self is part of the reduction.
+        # A first step is added to replace the impacted value by another one
+        # chosen in a way that the results of the reduction is not changed
+        # whether or not it takes part in it.
+        # It is -inf if the reduction is max, inf for min, 0 for add, 1 for mul.
+        # mean is not supported.
+        dtype = src.dtype or self.dtype
+        post_process_after_cast_like = False
+        if dtype is None:
+            dtype = ir.DataType.FLOAT
+            cast_like = True
+        else:
+            cast_like = False
+        # dtype should be not None.
+        if onnx_reduce == "max":
+            if dtype in {
+                ir.DataType.FLOAT16,
+                ir.DataType.FLOAT,
+                ir.DataType.DOUBLE,
+            }:
+                value = ir.tensor([np.finfo(dtype.numpy()).min], dtype=dtype)
+            elif dtype in {ir.DataType.BFLOAT16}:
+                value = ir.tensor([ml_dtypes.finfo(dtype.numpy()).min], dtype=dtype)
+            else:
+                value = ir.tensor([np.iinfo(dtype.numpy()).min], dtype=dtype)
+            reduction_init = "min"
+        elif onnx_reduce == "min":
+            post_process_after_cast_like = cast_like
+            if dtype in {
+                ir.DataType.FLOAT16,
+                ir.DataType.FLOAT,
+                ir.DataType.DOUBLE,
+            }:
+                value = ir.tensor([np.finfo(dtype.numpy()).max], dtype=dtype)
+            elif dtype == ir.DataType.BFLOAT16:
+                value = ir.tensor([ml_dtypes.finfo(dtype.numpy()).max], dtype=dtype)
+            else:
+                # Cast 1e20 into int32 returns the min value -2147483648
+                value = ir.tensor([np.iinfo(dtype.numpy()).max], dtype=dtype)
+            reduction_init = "max"
+        elif onnx_reduce == "add":
+            value = ir.tensor([0], dtype=dtype)
+            reduction_init = "none"
+        elif onnx_reduce == "mul":
+            value = ir.tensor([1], dtype=dtype)
+            reduction_init = "none"
+        else:
+            value = 0
+            reduction_init = "none"
+
+        cst = op.ConstantOfShape(op.Shape(src), value=value)
+        if cast_like:
+            cst = op.CastLike(cst, self)
+            if post_process_after_cast_like:
+                cst = op.Max(cst, op.Neg(op.Add(cst, op.CastLike(1, cst))))
+        self = op.ScatterElements(self, index, cst, axis=dim, reduction=reduction_init)
+
     result = op.ScatterElements(self, index, src, axis=dim, reduction=onnx_reduce)
+
     if self_is_scalar:
         result = op.Squeeze(result)
     return result

onnxscript/ir/_enums.py

@@ -74,6 +74,8 @@ def from_numpy(cls, dtype: np.dtype) -> DataType:
             TypeError: If the data type is not supported by ONNX.
         """
         if dtype not in _NP_TYPE_TO_DATA_TYPE:
+            if dtype == ml_dtypes.bfloat16 or (hasattr(dtype, "names") and dtype.names == ("bfloat16",)):
+                return DataType.BFLOAT16
             raise TypeError(f"Unsupported numpy data type: {dtype}")
         return cls(_NP_TYPE_TO_DATA_TYPE[dtype])
 

onnxscript/optimizer/_constant_folding.py

@@ -827,7 +827,11 @@ def _do_inference(self, node: ir.Node) -> None:
         def get_constant_value(x: ir.Value) -> onnx.TensorProto | None:
             value = _get_numpy_value(x)
             if isinstance(value, np.ndarray) and value.size < 20:
-                return onnx.numpy_helper.from_array(value, x.name)
+                try:
+                    return onnx.numpy_helper.from_array(value, x.name)
+                except ValueError:
+                    # This happens for bfloat16 and old versions of ONNX.
+                    return None
             return None
 
         def get_type(value: ir.Value) -> onnx.TypeProto | None:

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# TODO(pytorch/pytorch#129279): Migrate these tests to the PyTorch repo
+
+import itertools
+import unittest
+
+import onnxruntime
+import torch
+
+from tests.common import testutils
+
+
+class TorchLibe2eTest(testutils.TestBase):
+    def test_investigate_one_particular_model(self):
+        """This test can be used to investigate a particular issue."""
+        red, include, stype = "amin", False, "int32"
+        dtype = getattr(torch, stype)
+
+        class Model(torch.nn.Module):
+            def __init__(self, include, red):
+                super().__init__()
+                self.include = include
+                self.red = red
+
+            def forward(self, x, indices, updates):
+                x = x.clone()
+                return x.scatter_reduce(
+                    0, indices, updates, self.red, include_self=self.include
+                )
+
+        model = Model(include, red)
+        xs = (
+            torch.tensor([[-2, 0, 2], [2, -2, 0]], dtype=dtype),
+            torch.tensor([[0, 0, 0], [1, 1, 1]], dtype=torch.int64),
+            torch.tensor([[-1, -1, -1], [-1, -1, -1]], dtype=dtype),
+        )
+        expected = model(*xs)
+        model_path = (
+            f"test_aten_scatter_{red}_"
+            f"{'include' if include else 'exclude'}_{stype}.onnx"
+        )
+        torch.onnx.export(model, xs, model_path, dynamo=True)
+        feeds = dict(zip(["x", "indices", "updates"], [x.numpy() for x in xs]))
+
+        sess_options = onnxruntime.SessionOptions()
+        sess = onnxruntime.InferenceSession(
+            model_path, sess_options=sess_options, providers=["CPUExecutionProvider"]
+        )
+        got = sess.run(None, feeds)[0]
+        torch.testing.assert_close(
+            expected, torch.from_numpy(got), atol=1e-5, rtol=1e-5
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/function_libs/torch_lib/ops_test_data.py

@@ -2025,26 +2025,30 @@ def _where_input_wrangler(
         variant_name="mean",
         reason="ONNX doesn't support reduce='mean' option",
     )
-    .skip(
-        # ONNX has not include_self parameter and default is include_self=True mode
-        matcher=lambda sample: sample.kwargs.get("include_self") is False,
-        reason="ONNX does't support include_self=False option",
+    .xfail(
+        variant_name="prod",
+        dtypes=(torch.float16, torch.float64),
+        reason="fixme: MLFloat16 data type is not supported with ScatterElements opset 16 when reduction is 'mul'",
     )
     .xfail(
-        variant_name="amax",
-        reason="fixme: MLFloat16 data type is not supported with ScatterElements opset 18 when reduction is 'max'",
+        variant_name="sum",
+        dtypes=(torch.float16, torch.float64),
+        reason="fixme: MLFloat16 data type is not supported with ScatterElements opset 18 when reduction is 'add'",
     )
     .xfail(
-        variant_name="amin",
-        reason="fixme: MLFloat16 data type is not supported with ScatterElements opset 18 when reduction is 'min'",
+        variant_name="mean",
+        dtypes=(torch.bfloat16,),
+        reason="onnxruntime does not support ml_dtypes.bfloat16",
     )
     .xfail(
         variant_name="prod",
-        reason="fixme: MLFloat16 data type is not supported with ScatterElements opset 18 when reduction is 'prod'",
+        dtypes=(torch.bfloat16,),
+        reason="onnxruntime does not support ml_dtypes.bfloat16",
     )
     .xfail(
         variant_name="sum",
-        reason="fixme: MLFloat16 data type is not supported with ScatterElements opset 18 when reduction is 'add'",
+        dtypes=(torch.bfloat16,),
+        reason="onnxruntime does not support ml_dtypes.bfloat16",
     ),
     TorchLibOpInfo("ops.aten.slice_scatter", core_ops.aten_slice_scatter),
     TorchLibOpInfo("slice", core_ops.aten_slice),