[Relay][Pytorch] Add support for aten::bitwise_and (#16105)

add support for aten::bitwise_and

python/tvm/relay/frontend/pytorch.py

@@ -2327,6 +2327,14 @@ def bitwise_xor(self, inputs, input_types):
 
         return _op.bitwise_xor(lhs, rhs)
 
+    def bitwise_and(self, inputs, input_types):
+        lhs = inputs[0]
+        rhs = inputs[1]
+        lhs = _op.cast(lhs, "bool") if input_types[0] == "bool" else _op.cast(lhs, "int")
+        rhs = _op.cast(rhs, "bool") if input_types[1] == "bool" else _op.cast(rhs, "int")
+
+        return _op.bitwise_and(lhs, rhs)
+
     def logical_not(self, inputs, input_types):
         data = _wrap_const(inputs[0])
         return _op.logical_not(_op.cast(data, "bool"))
@@ -4033,6 +4041,7 @@ def create_convert_map(self):
             "aten::logical_xor": self.logical_xor,
             "aten::bitwise_not": self.bitwise_not,
             "aten::bitwise_xor": self.bitwise_xor,
+            "aten::bitwise_and": self.bitwise_and,
             "aten::Bool": self.Bool,
             "aten::Float": self.Float,
             "aten::rsub": self.rsub,

tests/python/frontend/pytorch/test_forward.py

@@ -3695,6 +3695,33 @@ def forward(self, *args):
     verify_model(BitwiseXor2().float().eval(), input_data=[lhs])
 
 
+def test_forward_bitwise_and():
+    """test_forward_bitwise_and"""
+    torch.set_grad_enabled(False)
+
+    class BitwiseAnd1(Module):
+        def forward(self, *args):
+            return torch.bitwise_and(args[0], args[1])
+
+    class BitwiseAnd2(Module):
+        def forward(self, *args):
+            rhs = torch.tensor([1, 0, 3], dtype=torch.int8)
+            if torch.cuda.is_available():
+                rhs = rhs.cuda()
+            return torch.bitwise_and(args[0], rhs)
+
+    lhs = torch.tensor([-1, -2, 3], dtype=torch.int8)
+    rhs = torch.tensor([1, 0, 3], dtype=torch.int8)
+    verify_model(BitwiseAnd1().float().eval(), input_data=[lhs, rhs])
+
+    lhs = torch.tensor([True, True, False])
+    rhs = torch.tensor([False, True, False])
+    verify_model(BitwiseAnd1().float().eval(), input_data=[lhs, rhs])
+
+    lhs = torch.tensor([-1, -2, 3], dtype=torch.int8)
+    verify_model(BitwiseAnd2().float().eval(), input_data=[lhs])
+
+
 @tvm.testing.uses_gpu
 def test_forward_logical_xor():
     """test_forward_logical_xor"""