add equal

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -84,6 +84,35 @@ def test_op_binary(self):
             else:
                 self.assertEmpty(v.value)
 
+    def test_op_binary_cmp(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Greater", ["X", "Y"], ["a"]),
+                    oh.make_node("GreaterOrEqual", ["X", "Y"], ["b"]),
+                    oh.make_node("Less", ["X", "Y"], ["c"]),
+                    oh.make_node("LessOrEqual", ["X", "Y"], ["d"]),
+                    oh.make_node("And", ["a", "b"], ["ab"]),
+                    oh.make_node("Or", ["c", "d"], ["cd"]),
+                    oh.make_node("And", ["ab", "cd"], ["Z"]),
+                ],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["a", "b"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b"]),
+                ],
+                [oh.make_tensor_value_info("Z", onnx.TensorProto.BOOL, ["a", "b"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        onnx.checker.check_model(model)
+        self._finalize_test(
+            model,
+            torch.abs(torch.rand(3, 4, 5, dtype=torch.float32)),
+            torch.abs(torch.rand(3, 4, 5, dtype=torch.float32)),
+        )
+
     def test_op_slice_squeeze(self):
         X = oh.make_tensor_value_info("X", TFLOAT, [None, None])
         starts = oh.make_tensor_value_info("starts", TINT64, [None])
@@ -220,6 +249,7 @@ def test_op_matmul(self):
             model,
             torch.rand((4, 5, 6, 7), dtype=torch.float32),
             torch.rand((4, 5, 7, 11), dtype=torch.float32),
+            atol=1e-6,
         )
 
     def test_op_unsqueeze(self):

onnx_diagnostic/ext_test_case.py

@@ -907,11 +907,13 @@ def assertEqualArray(
         except AssertionError as e:
             expected_max = numpy.abs(expected).max()
             expected_value = numpy.abs(value).max()
+            te = expected.astype(int) if expected.dtype == numpy.bool_ else expected
+            tv = value.astype(int) if value.dtype == numpy.bool_ else value
             rows = [
                 f"{msg}\n{e}" if msg else str(e),
                 f"expected max value={expected_max}",
                 f"expected computed value={expected_value}\n",
-                f"ratio={expected / value}\ndiff={expected - value}",
+                f"ratio={te / tv}\ndiff={te - tv}",
             ]
             raise AssertionError("\n".join(rows))  # noqa: B904
 

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -1,6 +1,18 @@
 from ._op_run import OpRun, OpRunValue
 from .access_ops import Gather_1, Slice_13
-from .binary_ops import Add_1, Div_1, MatMul_1, Mul_1, Sub_1
+from .binary_ops import (
+    And_1,
+    Add_1,
+    Div_1,
+    Greater_1,
+    GreaterOrEqual_1,
+    Less_1,
+    LessOrEqual_1,
+    MatMul_1,
+    Mul_1,
+    Or_1,
+    Sub_1,
+)
 from .other_ops import Cast_6, Concat_1, Transpose_1
 from .nn_ops import Softmax_13, Tanh_6
 from .shape_ops import Reshape_14, Shape_15, Squeeze_13, Unsqueeze_13

onnx_diagnostic/reference/torch_ops/binary_ops.py

@@ -1,6 +1,13 @@
 from . import OpRun, OpRunValue
 
 
+class And_1(OpRun):
+    """And"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor & y.tensor)
+
+
 class Add_1(OpRun):
     """Add"""
 
@@ -15,6 +22,41 @@ def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
         return OpRunValue(x.tensor / y.tensor)
 
 
+class Equal_1(OpRun):
+    """Equal"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor == y.tensor)
+
+
+class Greater_1(OpRun):
+    """Greater"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor > y.tensor)
+
+
+class GreaterOrEqual_1(OpRun):
+    """GreaterOrEqual"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor >= y.tensor)
+
+
+class Less_1(OpRun):
+    """Less"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor < y.tensor)
+
+
+class LessOrEqual_1(OpRun):
+    """LessOrEqual"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor <= y.tensor)
+
+
 class MatMul_1(OpRun):
     """MatMul"""
 
@@ -29,6 +71,13 @@ def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
         return OpRunValue(x.tensor * y.tensor)
 
 
+class Or_1(OpRun):
+    """Or"""
+
+    def run(self, x: OpRunValue, y: OpRunValue) -> OpRunValue:
+        return OpRunValue(x.tensor | y.tensor)
+
+
 class Sub_1(OpRun):
     """Sub"""