Enable torch.is_complex in Dynamo tracing (#103154)

fduwjj · fduwjj · commit 1af8ff4570c6 · 2023-06-08T01:26:35.000Z
Pull Request resolved: #103154 Approved by: https://github.com/yanboliang ghstack-source-id: 6b16fc3
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
@@ -436,6 +436,23 @@ def test_get_default_dtype(x):
         else:
             return x - 1
 
+    @make_test
+    def test_get_autocast_gpu_dtype(x):
+        dtype = torch.get_autocast_gpu_dtype()
+        return x.type(dtype)
+
+    @make_test
+    def test_get_calculate_correct_fan(x):
+        fan_in = torch.nn.init._calculate_correct_fan(x, "fan_in")
+        return x + fan_in
+
+    @make_test
+    def test_is_complex(x):
+        if torch.is_complex(x):
+            return x + 1
+        else:
+            return x - 1
+
     @make_test
     def test_device(x):
         if not x.is_cuda:
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
@@ -11,9 +11,14 @@
 import torch.fx
 import torch.nn
 import torch.onnx.operators
+<<<<<<< HEAD
 from torch._dynamo.utils import get_fake_value, get_real_value
 from torch._dynamo.variables import SymNodeVariable
 from torch._dynamo.variables.user_defined import ProcessGroupVariable
+=======
+from torch._dynamo.utils import get_fake_value, get_real_value, torch_np
+from torch._dynamo.variables import SymNodeVariable, UserFunctionVariable
+>>>>>>> fecd5f75277... Enable torch.nn.init._calculate_correct_fan in dynamo tracing
 from torch._guards import GuardsCheckpointState, Source
 from torch.utils import _pytree as pytree
 
@@ -73,6 +78,7 @@
     torch.is_autocast_cache_enabled,
     torch.is_autocast_cpu_enabled,
     torch.is_autocast_enabled,
+    torch.is_complex,
     torch.is_floating_point,
     torch.nn.functional._Reduction.get_enum,
 ]
@@ -528,6 +534,10 @@ def get_state_from_generator():
             assert len(args) == 1, "Expected one arg (pg)"
             assert isinstance(args[0], ProcessGroupVariable)
             return ConstantVariable(self.value(args[0].as_python_constant()))
+        elif self.value == torch.nn.init._calculate_correct_fan:
+            return UserFunctionVariable(
+                torch.nn.init._calculate_correct_fan, **options
+            ).call_function(tx, args, {})
         else:
             any_symints_or_symfloats = any(isinstance(x, SymNodeVariable) for x in args)
             all_ints_or_floats = all(
