update

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -435,6 +435,23 @@ def test_op_reduce_min(self):
             atol=1e-6,
         )
 
+    def test_op_reduce_min_17(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("ReduceMin", ["X"], ["Z"], axes=[1])],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 17)],
+        )
+        self._finalize_test(
+            model,
+            torch.rand(3, 4, 5, dtype=torch.float32),
+            atol=1e-6,
+        )
+
     def test_op_reduce_sum(self):
         model = oh.make_model(
             oh.make_graph(
@@ -1104,6 +1121,21 @@ def test_conv(self):
         B = torch.tensor([[[[0]]]], dtype=torch.float32)
         self._finalize_test(model, X, W, B, use_ort=True)
 
+    def test_nonzero(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("NonZero", ["X"], ["Y"])],
+                "g",
+                [oh.make_tensor_value_info("X", TFLOAT, [None, None])],
+                [oh.make_tensor_value_info("Y", TINT64, [None, None])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+
+        self._finalize_test(
+            model, torch.tensor([[1, 0], [1, 1]], dtype=torch.float32), use_ort=True
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -18,8 +18,16 @@
 from .controlflow_ops import If_1, Loop_16
 from .generator_ops import Range_11
 from .nn_ops import Conv_11, LayerNormalization_17, Softmax_13, Tanh_6
-from .other_ops import Cast_6, CastLike_15, Concat_1, Transpose_1, Trilu_14, Where_9
-from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_18
+from .other_ops import (
+    Cast_6,
+    CastLike_15,
+    NonZero_13,
+    Concat_1,
+    Transpose_1,
+    Trilu_14,
+    Where_9,
+)
+from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_17, ReduceMin_18, ReduceSum_13
 from .sequence_ops import ConcatFromSequence_11, SequenceEmpty_11, SequenceInsert_11
 from .shape_ops import (
     ConstantOfShape_9,

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -45,6 +45,13 @@ def run(self, *data: OpRunTensor) -> OpRunTensor:
         return OpRunTensor(torch.cat([t.tensor for t in data], axis=self.axis))
 
 
+class NonZero_13(OpRun):
+    "NonZero"
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.nonzero(x.tensor).T)
+
+
 class Transpose_1(OpRun):
     "Transpose"
 

onnx_diagnostic/reference/torch_ops/reduce_ops.py

@@ -27,6 +27,12 @@ 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 ReduceOpAxes(ReduceOp):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.axes = self.get_attribute_ints(node, "axes", [])
+
+
 class ReduceMax_18(ReduceOp):
     """ReduceMax"""
 
@@ -65,6 +71,27 @@ def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor
         return OpRunTensor(t)
 
 
+class ReduceMin_17(ReduceOpAxes):
+    """ReduceMin"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        assert self.stash_type is None, f"Not implemented with stash_type={self.stash_type}"
+        axes = self.axes
+        if not axes:
+            assert (
+                not self.keepdims
+            ), f"axes is Empty, keepdims={self.keepdims} for {self.__class__.__name__}"
+            return OpRunTensor(torch.min(x.tensor).values)
+        taxes = tuple(axes)
+        if len(taxes) == 1:
+            t = x.tensor.min(taxes[0], keepdim=self.keepdims)
+            return OpRunTensor(t.values)
+        t = x.tensor
+        for a in reversed(taxes):
+            t = t.min(a, keepdim=self.keepdims).values
+        return OpRunTensor(t)
+
+
 class ReduceMin_18(ReduceOp):
     """ReduceMin"""
 
@@ -85,7 +112,7 @@ def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor
         return OpRunTensor(t)
 
 
-class ReduceSum_18(ReduceOp):
+class ReduceSum_13(ReduceOp):
     """ReduceSum"""
 
     def run(self, x: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor: