Add BOOL variants to bitwise ops | fix(torchlib) (#907)

In torchbench we see errors like `[ONNXRuntimeError] : 10 :
INVALID_GRAPH : Load model from
bench_dynamo_onnx_model/attention_is_all_you_need_pytorch/model.onnx
failed:This is an invalid model. Type Error: Type 'tensor(bool)' of
input parameter (unsqueeze_1) of operator (BitwiseAnd) in node (n0__9)
is invalid.`

This change fixes them by registering the Boolean variants from the
exiting logical* implementations.

The reason for not creating new bitwise* variants is that the
implementation is exactly the same as the logical* ones, which are
covered by test cases. The bool bitwise* BOOL variants, if implemented
separately, besides completely identical logic, would require test case
duplication, which imo is not worth it.

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -1007,9 +1007,17 @@ def aten_binomial(
     raise NotImplementedError()
 
 
-@torch_op("aten::bitwise_and")
+@torch_op(
+    (
+        "aten::bitwise_and",
+        "aten::bitwise_and.Tensor",
+        "aten::bitwise_and.Scalar",
+        "aten::bitwise_and.Scalar_Tensor",
+    )
+)
 def aten_bitwise_and(self: TInt, other: TInt) -> TInt:
     """bitwise_and.Tensor(Tensor self, Tensor other) -> Tensor"""
+    # logical_and implements the BOOL variant
 
     return op.BitwiseAnd(self, other)
 
@@ -1024,19 +1032,22 @@ def aten_bitwise_left_shift(self: TInt, other: TInt) -> TInt:
 @torch_op("aten::bitwise_not")
 def aten_bitwise_not(self: TInt) -> TInt:
     """bitwise_not(Tensor self) -> Tensor"""
+    # logical_not implements the BOOL variant
 
     return op.BitwiseNot(self)
 
 
-@torch_op("aten::bitwise_not")
-def aten_bitwise_not_bool(self: BOOL) -> BOOL:
-    """bitwise_not(Tensor self) -> Tensor"""
-    return op.Not(self)
-
-
-@torch_op("aten::bitwise_or")
+@torch_op(
+    (
+        "aten::bitwise_or",
+        "aten::bitwise_or.Tensor",
+        "aten::bitwise_or.Scalar",
+        "aten::bitwise_or.Scalar_Tensor",
+    )
+)
 def aten_bitwise_or(self: TInt, other: TInt) -> TInt:
     """bitwise_or.Tensor(Tensor self, Tensor other) -> Tensor"""
+    # logical_or implements the BOOL variant
 
     return op.BitwiseOr(self, other)
 
@@ -1048,9 +1059,17 @@ def aten_bitwise_right_shift(self: TInt, other: TInt) -> TInt:
     return op.BitShift(self, other, direction="RIGHT")
 
 
-@torch_op("aten::bitwise_xor")
+@torch_op(
+    (
+        "aten::bitwise_xor",
+        "aten::bitwise_xor.Tensor",
+        "aten::bitwise_xor.Scalar",
+        "aten::bitwise_xor.Scalar_Tensor",
+    )
+)
 def aten_bitwise_xor(self: TInt, other: TInt) -> TInt:
     """bitwise_xor.Tensor(Tensor self, Tensor other) -> Tensor"""
+    # logical_xor implements the BOOL variant
 
     return op.BitwiseXor(self, other)
 
@@ -3734,28 +3753,52 @@ def aten_logdet(self: TFloat) -> TFloat:
     return op.Log(op.Det(self))
 
 
-@torch_op("aten::logical_and")
+@torch_op(
+    (
+        "aten::logical_and",
+        "aten::bitwise_and",
+        "aten::bitwise_and.Tensor",
+        "aten::bitwise_and.Scalar",
+        "aten::bitwise_and.Scalar_Tensor",
+    )
+)
 def aten_logical_and(self: BOOL, other: BOOL) -> BOOL:
     """logical_and(Tensor self, Tensor other) -> Tensor"""
 
     return op.And(self, other)
 
 
-@torch_op("aten::logical_not")
+@torch_op(("aten::logical_not", "aten::bitwise_not"))
 def aten_logical_not(self: BOOL) -> BOOL:
     """logical_not(Tensor self) -> Tensor"""
 
     return op.Not(self)
 
 
-@torch_op("aten::logical_or")
+@torch_op(
+    (
+        "aten::logical_or",
+        "aten::bitwise_or",
+        "aten::bitwise_or.Tensor",
+        "aten::bitwise_or.Scalar",
+        "aten::bitwise_or.Scalar_Tensor",
+    )
+)
 def aten_logical_or(self: BOOL, other: BOOL) -> BOOL:
     """logical_or(Tensor self, Tensor other) -> Tensor"""
 
     return op.Or(self, other)
 
 
-@torch_op("aten::logical_xor")
+@torch_op(
+    (
+        "aten::logical_xor",
+        "aten::bitwise_xor",
+        "aten::bitwise_xor.Tensor",
+        "aten::bitwise_xor.Scalar",
+        "aten::bitwise_xor.Scalar_Tensor",
+    )
+)
 def aten_logical_xor(self: BOOL, other: BOOL) -> BOOL:
     """logical_xor(Tensor self, Tensor other) -> Tensor"""