reducemin

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -345,6 +345,33 @@ def test_op_tanh(self):
             model, torch.abs(torch.rand(3, 4, 5, dtype=torch.float32)), atol=1e-6
         )
 
+    def test_op_reduce_min(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceMin", ["X", "axes"], ["Z"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"]),
+                    oh.make_tensor_value_info("axes", TINT64, ["f"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            torch.tensor([1], dtype=torch.int64),
+            atol=1e-6,
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            torch.tensor([1, 2], dtype=torch.int64),
+            atol=1e-6,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -15,5 +15,6 @@
 )
 from .other_ops import Cast_6, Concat_1, Transpose_1
 from .nn_ops import Softmax_13, Tanh_6
+from .reduce_ops import ReduceMin_18
 from .shape_ops import Reshape_14, Shape_15, Squeeze_13, Unsqueeze_13
 from .unary_ops import Neg_1, Not_1, Reciprocal_1

onnx_diagnostic/reference/torch_ops/reduce_ops.py

@@ -0,0 +1,41 @@
+from typing import Optional
+import onnx
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRun, OpRunValue
+
+
+class ReduceOp(OpRun):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.keepdims = bool(self.get_attribute_int(node, "keepdims", 1))
+        self.noop_with_empty_axes = bool(
+            self.get_attribute_int(node, "noop_with_empty_axes", 0)
+        )
+        assert isinstance(
+            self.keepdims, bool
+        ), f"Unexpected value for attribute keepdims={self.keepdims!r}"
+        assert isinstance(self.noop_with_empty_axes, bool), (
+            f"Unexpected value for attribute "
+            f"noop_with_empty_axes={self.noop_with_empty_axes!r}"
+        )
+        assert (
+            not self.noop_with_empty_axes
+        ), f"Not implemented with noop_with_empty_axes={self.noop_with_empty_axes}"
+        # this is out of spec
+        stash_type = self.get_attribute_int(node, "stash_type", None)
+        self.stash_type = None if stash_type is None else onnx_dtype_to_torch_dtype(stash_type)
+
+
+class ReduceMin_18(ReduceOp):
+    """ReduceMin"""
+
+    def run(self, x: OpRunValue, axes: OpRunValue) -> OpRunValue:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        taxes = axes.as_tuple_int
+        if len(taxes) == 1:
+            t = x.tensor.min(taxes[0], keepdim=self.keepdims)
+            return OpRunValue(t.values)
+        t = x.tensor
+        for a in reversed(taxes):
+            t = t.min(a, keepdim=self.keepdims).values
+        return OpRunValue(t)