[Dynamo] Supports torch._C._is_any_autocast_enabled (#126196)

Fixes #126026 Pull Request resolved: #126196 Approved by: https://github.com/anijain2305
pytorch · May 19, 2024 · 976f0f2 · 976f0f2
1 parent a745003
commit 976f0f2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/test/dynamo/test_functions.py b/test/dynamo/test_functions.py
@@ -642,6 +642,13 @@ def test_get_autocast_gpu_dtype(x):
         dtype = torch.get_autocast_gpu_dtype()
         return x.type(dtype)
 
+    @make_test
+    def test_is_any_autocast_enabled(x):
+        if torch._C._is_any_autocast_enabled():
+            return x + 1
+        else:
+            return x - 1
+
     @make_test
     def test_list_compare_polyfill(x):
         for a, b, c in [

diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
@@ -83,6 +83,7 @@
     torch._assert,
     torch._utils._get_device_index,
     torch._C._get_cublas_allow_tf32,
+    torch._C._is_any_autocast_enabled,
     torch.cuda.get_device_properties,
     torch.cuda.is_available,
     torch.distributed.is_available,