[ONNX] Export logical_not (pytorch#96315)

Fixes pytorch#95154

Pull Request resolved: pytorch#96315
Approved by: https://github.com/justinchuby

Author: BowenBao

Diff for: test/onnx/test_op_consistency.py

@@ -305,6 +305,7 @@ def reason_flaky() -> str:
 TESTED_OPS: frozenset[str] = frozenset(
     [
         "ceil",
+        "logical_not",
         "sqrt",
         "stft",
         "t",

Diff for: test/onnx/test_pytorch_onnx_onnxruntime.py

@@ -4920,6 +4920,24 @@ def forward(self, x, y):
         y = torch.randint(10, (2, 3, 5), dtype=torch.long)
         self.run_test(XorModel(), input_args=(x, y))
 
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_logical_not(self):
+        class NotModel(torch.nn.Module):
+            def forward(self, x):
+                return torch.logical_not(x)
+
+        x = torch.randint(0, 2, (5, 5), dtype=torch.bool)
+        self.run_test(NotModel(), input_args=(x,))
+
+        x = torch.randint(10, (5, 5), dtype=torch.int32)
+        self.run_test(NotModel(), input_args=(x,))
+
+        x = torch.randint(10, (5, 5), dtype=torch.double)
+        self.run_test(NotModel(), input_args=(x,))
+
+        x = torch.randint(10, (2, 3, 5), dtype=torch.float32)
+        self.run_test(NotModel(), input_args=(x,))
+
     @skipIfUnsupportedMinOpsetVersion(11)  # float equal added after opset 11
     def test_eq(self):
         class EqualModel(torch.nn.Module):

Diff for: torch/onnx/symbolic_opset9.py

@@ -140,6 +140,7 @@
     "log1p",
     "log2",
     "logical_and",
+    "logical_not",
     "logical_or",
     "logical_xor",
     "logsumexp",
@@ -2295,6 +2296,12 @@ def logical_xor(g: jit_utils.GraphContext, input, other):
     return g.op("Xor", input, other)
 
 
+@_onnx_symbolic("aten::logical_not")
+@_beartype.beartype
+def logical_not(g: jit_utils.GraphContext, input):
+    return g.op("Not", g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.BOOL))
+
+
 @_onnx_symbolic("aten::__rshift_")
 @_beartype.beartype
 def __rshift_(g: jit_utils.GraphContext, self, other):