Delete torch.autograd.function.traceable APIs

zou3519 · zou3519 · commit 464ab7064b69 · 2024-03-27T11:59:28.000-07:00
We deprecated them in 2.3 with plans to delete in 2.4. Very few OSS
repos use this flag at all and it also does nothing.

[ghstack-poisoned]
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -510,47 +510,6 @@ def test_not_implemented_fwad(self):
                 # if forward AD ends up being implemented for torch.igamma, choose a different op
                 torch.igamma(dual_x, dual_x)
 
-    def test_traceable_deprecated(self):
-        class MyFunction(Function):
-            @staticmethod
-            def forward(ctx, x):
-                return x * 2
-
-            @staticmethod
-            def backward(ctx, gO):
-                return gO * 2
-
-        with self.assertWarnsRegex(UserWarning, "is_traceable .*is deprecated"):
-            MyFunction.is_traceable
-
-        with self.assertWarnsRegex(UserWarning, "is_traceable .*is deprecated"):
-            MyFunction.is_traceable = True
-
-        with self.assertWarnsRegex(UserWarning, "is_traceable .*is deprecated"):
-            class MyFunction(Function):
-                is_traceable = True
-
-                @staticmethod
-                def forward(ctx, x):
-                    return x * 2
-
-                @staticmethod
-                def backward(ctx, gO):
-                    return gO * 2
-
-        with self.assertWarnsRegex(UserWarning, "traceable .*is deprecated"):
-            @torch.autograd.function.traceable
-            class MyFunction(Function):
-                is_traceable = True
-
-                @staticmethod
-                def forward(ctx, x):
-                    return x * 2
-
-                @staticmethod
-                def backward(ctx, gO):
-                    return gO * 2
-
     def test_will_engine_execute_node(self):
         counter = [0]
 
diff --git a/torch/autograd/function.py b/torch/autograd/function.py
@@ -18,7 +18,6 @@
     "FunctionMeta",
     "Function",
     "once_differentiable",
-    "traceable",
     "InplaceFunction",
     "NestedIOFunction",
 ]
@@ -311,23 +310,6 @@ def _compiled_autograd_key(self):
         return self._forward_cls._compiled_autograd_key(self)  # type: ignore[attr-defined]
 
 
-def _warn_traceable_deprecated():
-    warnings.warn(
-        "The is_traceable field on torch.autograd.Function is deprecated "
-        "and will be removed in PyTorch 2.4.",
-        stacklevel=3,
-    )
-
-
-class _IsTraceableWarning:
-    def __init__(self, value):
-        self.value = value
-
-    def __get__(self, obj, klass=None):
-        _warn_traceable_deprecated()
-        return self.value
-
-
 class FunctionMeta(type):
     """Function metaclass.
 
@@ -347,24 +329,8 @@ def __init__(cls, name, bases, attrs):
         )
         cls._backward_cls = backward_fn
 
-        if "is_traceable" in attrs:
-            if attrs["is_traceable"] is True:
-                _warn_traceable_deprecated()
-                # already emitted warning, no need to install _IsTraceableWarning
-            else:
-                cls.is_traceable = _IsTraceableWarning(attrs.pop("is_traceable"))
-
         super().__init__(name, bases, attrs)
 
-    def __setattr__(cls, name, value):
-        if name == "is_traceable" and value is True:
-            warnings.warn(
-                "The is_traceable field on torch.autograd.Function is deprecated "
-                "and will be removed in PyTorch 2.4.",
-                stacklevel=2,
-            )
-        return super().__setattr__(name, value)
-
 
 class _SingleLevelFunction(
     _C._FunctionBase, FunctionCtx, _HookMixin, metaclass=FunctionMeta
@@ -540,9 +506,6 @@ def __call__(self, *args, **kwargs):
             "(Example: https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)"
         )
 
-    # for the tracer
-    is_traceable = False
-
     """
     Bool that specifies if PyTorch should attempt to autogenerate
     :func:`torch.vmap` support for this autograd.Function. You may set this to
@@ -667,26 +630,6 @@ def fake_requires_grad(var):
     return wrapper
 
 
-def traceable(fn_cls):
-    r"""Mark Function as traceable for the JIT.
-
-    Traceable functions have additional restrictions - they can't pass any
-    data-dependent values to backward (e.g. Prod passes the output, which makes
-    it non-traceable), and their backward should be implemented entirely in terms
-    of operations on autograd Tensors in all cases.
-
-    DON'T USE THIS DECORATOR. IT IS FOR INTERNAL USE ONLY AND SHOULD BE HANDLED WITH
-    CARE (or can give incorrect results otherwise).
-    """
-    warnings.warn(
-        "torch.autograd.function.traceable is deprecated "
-        "and will be removed in PyTorch 2.4.",
-        stacklevel=2,
-    )
-    fn_cls.is_traceable = True
-    return fn_cls
-
-
 class InplaceFunction(Function):
     r"""
     This class is here only for backward compatibility reasons.
