reduce

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -345,6 +345,60 @@ def test_op_tanh(self):
             model, torch.abs(torch.rand(3, 4, 5, dtype=torch.float32)), atol=1e-6
         )
 
+    def test_op_reduce_max(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceMax", ["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,
+        )
+
+    def test_op_reduce_mean(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceMean", ["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,
+        )
+
     def test_op_reduce_min(self):
         model = oh.make_model(
             oh.make_graph(
@@ -372,6 +426,33 @@ def test_op_reduce_min(self):
             atol=1e-6,
         )
 
+    def test_op_reduce_sum(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceSum", ["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,6 +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 .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_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

@@ -26,6 +26,34 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         self.stash_type = None if stash_type is None else onnx_dtype_to_torch_dtype(stash_type)
 
 
+class ReduceMax_18(ReduceOp):
+    """ReduceMax"""
+
+    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.max(taxes[0], keepdim=self.keepdims)
+            return OpRunValue(t.values)
+        t = x.tensor
+        for a in reversed(taxes):
+            t = t.max(a, keepdim=self.keepdims).values
+        return OpRunValue(t)
+
+
+class ReduceMean_18(ReduceOp):
+    """ReduceMean"""
+
+    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.mean(taxes[0], keepdim=self.keepdims)
+            return OpRunValue(t)
+        t = x.tensor.mean(taxes, keepdim=self.keepdims)
+        return OpRunValue(t)
+
+
 class ReduceMin_18(ReduceOp):
     """ReduceMin"""
 
@@ -39,3 +67,16 @@ def run(self, x: OpRunValue, axes: OpRunValue) -> OpRunValue:
         for a in reversed(taxes):
             t = t.min(a, keepdim=self.keepdims).values
         return OpRunValue(t)
+
+
+class ReduceSum_18(ReduceOp):
+    """ReduceSum"""
+
+    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.sum(taxes[0], keepdim=self.keepdims)
+            return OpRunValue(t)
+        t = x.tensor.sum(taxes, keepdim=self.keepdims)
+        return OpRunValue(t)