[ONNX] Export logical_not

ghstack-source-id: 3434436ec2854702de1abb49552a5f24e8f2db85 Pull Request resolved: #96315
pytorch · Mar 8, 2023 · 443e966 · 443e966
1 parent f63cd41
commit 443e966
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -4905,6 +4905,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 --git a/torch/onnx/symbolic_opset9.py b/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):