Torch eq and ne ops supports bool type.

Author: fukatani

coremltools/converters/mil/frontend/torch/ops.py

@@ -472,14 +472,26 @@ def listconstruct(context, node):
 @register_torch_op
 def eq(context, node):
     inputs = _get_inputs(context, node, expected=2)
-    equal_to = mb.equal(x=inputs[0], y=inputs[1], name=node.name)
+    x = inputs[0]
+    y = inputs[1]
+    if is_bool(x.dtype):
+        x = mb.cast(x=x, dtype='int32')
+    if is_bool(y.dtype):
+        y = mb.cast(x=y, dtype='int32')
+    equal_to = mb.equal(x=x, y=y, name=node.name)
     context.add(equal_to)
 
 
 @register_torch_op
 def ne(context, node):
     inputs = _get_inputs(context, node, expected=2)
-    equal_to = mb.not_equal(x=inputs[0], y=inputs[1], name=node.name)
+    x = inputs[0]
+    y = inputs[1]
+    if is_bool(x.dtype):
+        x = mb.cast(x=x, dtype='int32')
+    if is_bool(y.dtype):
+        y = mb.cast(x=y, dtype='int32')
+    equal_to = mb.not_equal(x=x, y=y, name=node.name)
     context.add(equal_to)
 
 

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -3152,86 +3152,37 @@ def forward(self, x):
         )
 
 
-class TestLogicalAnd(TorchBaseTest):
+class TestBitWiseLogical(TorchBaseTest):
     @pytest.mark.parametrize(
-        "backend, x_y",
+        "backend, x_y, op_string",
         itertools.product(
             backends,
             [
                 ([True, False, True, False], [True, True, False, False]),
                 ([[True, False], [True, False]], [[True, True], [False, False]]),
+                ([[True, False], [True, False]], [[1, 0], [2, 1]]),
                 ([-1.5, 0.0, 1.0, 0.0], [0.1, 2.5, 0.0, 0.0]),
                 ([2, 0, -1, 0, 5], [1, 1, 0, 0, -5]),
             ],
-        ),
-    )
-    def test_logical_and(self, backend, x_y):
-        class TestNet(nn.Module):
-            def __init__(self):
-                super(TestNet, self).__init__()
-
-            def forward(self, x, y):
-                return torch.logical_and(x, y)
-
-        model = TestNet()
-        x = torch.tensor(x_y[0])
-        y = torch.tensor(x_y[1])
-        self.run_compare_torch([x, y], model, backend=backend, input_as_shape=False)
-
-
-class TestLogicalOr(TorchBaseTest):
-    @pytest.mark.parametrize(
-        "backend, x_y",
-        itertools.product(
-            backends,
             [
-                ([True, False, True, False], [True, True, False, False]),
-                ([[True, False], [True, False]], [[True, True], [False, False]]),
-                ([-1.5, 0.0, 1.0, 0.0], [0.1, 2.5, 0.0, 0.0]),
-                ([2, 0, -1, 0, 5], [1, 1, 0, 0, -5]),
+                "eq",
+                "ne",
+                "logical_and",
+                "logical_or",
+                "logical_xor",
             ],
         ),
     )
-    def test_logical_or(self, backend, x_y):
-        class TestNet(nn.Module):
-            def __init__(self):
-                super(TestNet, self).__init__()
-
-            def forward(self, x, y):
-                return torch.logical_or(x, y)
-
-        model = TestNet()
+    def test_bitwise_logical(self, backend, x_y, op_string):
+        if not contains_op(torch, op_string):
+            return
+        op_func = getattr(torch, op_string)
+        model = ModuleWrapper(function=op_func)
         x = torch.tensor(x_y[0])
         y = torch.tensor(x_y[1])
         self.run_compare_torch([x, y], model, backend=backend, input_as_shape=False)
 
 
-class TestLogicalXor(TorchBaseTest):
-    @pytest.mark.parametrize(
-        "backend, x_y",
-        itertools.product(
-            backends,
-            [
-                ([True, False, True, False], [True, True, False, False]),
-                ([[True, False], [True, False]], [[True, True], [False, False]]),
-                ([-1.5, 0.0, 1.0, 0.0], [0.1, 2.5, 0.0, 0.0]),
-                ([2, 0, -1, 0, 5], [1, 1, 0, 0, -5]),
-            ],
-        ),
-    )
-    def test_logical_xor(self, backend, x_y):
-        class TestNet(nn.Module):
-            def __init__(self):
-                super(TestNet, self).__init__()
-
-            def forward(self, x, y):
-                return torch.logical_xor(x, y)
-
-        model = TestNet()
-        x = torch.tensor(x_y[0])
-        y = torch.tensor(x_y[1])
-        self.run_compare_torch([x, y], model, backend=backend, input_as_shape=False)
-
 
 class TestWhere(TorchBaseTest):
     @pytest.mark.parametrize(