Update on "Add deprecation notice for allow_in_graph"

ydwu4 · ydwu4 · commit 12edfb47a8c7 · 2026-03-11T12:26:00.000-07:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo

[ghstack-poisoned]
diff --git a/test/dynamo/test_decorators.py b/test/dynamo/test_decorators.py
@@ -221,6 +221,11 @@ def fn(a):
         self.assertEqual(cnts.frame_count, 1)
         self.assertEqual(cnts.op_count, 5)
 
+    def test_allow_in_graph_deprecation_warning(self):
+        with self.assertWarnsRegex(FutureWarning, "nonstrict_trace"):
+            torch._dynamo.allow_in_graph(my_custom_function)
+        torch._dynamo.disallow_in_graph(my_custom_function)
+
     def test_allow_in_graph_no_id_reuse(self):
         cnts = torch._dynamo.testing.CompileCounter()
 
diff --git a/torch/_dynamo/decorators.py b/torch/_dynamo/decorators.py
@@ -4,13 +4,12 @@
 
 import functools
 import inspect
-import warnings
 import weakref
 from collections.abc import Callable
 from dataclasses import dataclass
 from types import TracebackType
 from typing import Any, overload, TYPE_CHECKING, TypeVar
-from typing_extensions import ParamSpec
+from typing_extensions import deprecated, ParamSpec
 
 import torch
 import torch.utils._pytree as pytree
@@ -179,6 +178,11 @@ def assume_constant_result(fn):  # type: ignore[no-untyped-def]
     return fn
 
 
+@deprecated(
+    "torch._dynamo.allow_in_graph is deprecated and will be removed in a future version. "
+    "Use torch._dynamo.nonstrict_trace instead.",
+    category=FutureWarning,
+)
 def allow_in_graph(fn):  # type: ignore[no-untyped-def]
     """
     Tells the compiler frontend (Dynamo) to skip symbolic introspection of the function
@@ -188,12 +192,6 @@ def allow_in_graph(fn):  # type: ignore[no-untyped-def]
 
     WARNING: this API can be a footgun, please read the documentation carefully.
     """
-    warnings.warn(
-        "torch._dynamo.allow_in_graph is deprecated and will be removed in a future version. "
-        "Use torch._dynamo.nonstrict_trace instead.",
-        FutureWarning,
-        stacklevel=2,
-    )
     if isinstance(fn, (list, tuple)):
         return [allow_in_graph(x) for x in fn]
     assert callable(fn), "allow_in_graph expects a callable"
