[dynamo] Add support for tensor's is_complex method (#124927)

YangQun1 · pytorchmergebot · commit 91d565da0c5c · 2024-04-26T18:28:14.000Z
This PR is to add support for tensor's is_complex method in dynamo. Take the following code as an example: ```python def test_tensor_is_complex(x): if x.is_complex(): return x + 1 else: return x - 1 ``` Before this fix, the is_complex() call will cause a graph break "torch.* op returned non-Tensor bool call_method is_complex". After this fix, the graph break can be avoided. Fixes #122692 Pull Request resolved: #124927 Approved by: https://github.com/ezyang
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
@@ -681,6 +681,13 @@ def test_is_complex(x):
         else:
             return x - 1
 
+    @make_test
+    def test_tensor_is_complex(x):
+        if x.is_complex():
+            return x + 1
+        else:
+            return x - 1
+
     @make_test
     def test_get_privateuse1_name(x):
         if torch._C._get_privateuse1_backend_name() == "privateuseone":
diff --git a/torch/_dynamo/variables/tensor.py b/torch/_dynamo/variables/tensor.py
@@ -524,6 +524,10 @@ def method_is_floating_point(self):
         if self.dtype is not None:
             return ConstantVariable.create(self.dtype.is_floating_point)
 
+    def method_is_complex(self):
+        if self.dtype is not None:
+            return ConstantVariable.create(self.dtype.is_complex)
+
     def method_is_contiguous(self, memory_format=None):
         memory_format = (
             memory_format.as_python_constant()
