[Dynamo] Clean up old torch function flag (#149711)

mlazos · pytorchmergebot · commit c312d8c5018f · 2025-04-21T21:33:58.000Z
This is tracked via `SymbolicTorchFunctionState` now. Pull Request resolved: #149711 Approved by: https://github.com/StrongerXi, https://github.com/anijain2305
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
@@ -422,8 +422,6 @@ def __init__(
         self.should_exit = False
         self.unspec_variable_map: dict[str, UnspecializedPythonVariable] = {}
 
-        # Note this returns true iff TF Mode and TF Subclasses are enabled
-        self.torch_function_enabled = torch._C._is_torch_function_enabled()
         # This returns false if TF Overall (both mode and subclass) is disabled OR that TF Mode stack is empty
         self.torch_function_mode_enabled = torch._C._is_torch_function_mode_enabled()
         # This records the initial torch function mode stack for guarding
@@ -697,16 +695,6 @@ def save_global_state(self, out=None):
             ),
         )
 
-        # TODO - Consider having a torch level API for torch_function_state. As
-        # of now, we create a ref cycle by passing the
-        # output.set_torch_function_state to
-        # output.tracing_context.global_context.global_state. In the interim,
-        # the problem can be solved by manually set
-        # output.tracing_context.global_context.global_state to None at cleanup.
-        global_state["torch_function_enabled"] = (
-            self.set_torch_function_state,
-            self.torch_function_enabled,
-        )
         global_state["grad_enabled"] = (torch.set_grad_enabled, torch.is_grad_enabled())
 
         global_state["autocast_enabled"] = (
@@ -1876,9 +1864,6 @@ def cleanup(self) -> None:
         self.unspec_variable_map.clear()
         self.backward_state.clear()
 
-    def set_torch_function_state(self, enabled: bool) -> None:
-        self.torch_function_enabled = enabled
-
     def add_graph_finalizer(
         self, register_finalizer: Callable[[fx.GraphModule], None]
     ) -> None:
diff --git a/torch/_dynamo/variables/ctx_manager.py b/torch/_dynamo/variables/ctx_manager.py
@@ -738,7 +738,6 @@ def _call_func(self, tx: "InstructionTranslator", values):
         tx.symbolic_torch_function_state.torch_function_subclass_enabled = False
         if not self.only_subclass:
             tx.symbolic_torch_function_state.torch_function_mode_enabled = False
-        tx.output.set_torch_function_state(False)
 
     def module_name(self):
         return "torch._C"
diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
@@ -297,14 +297,11 @@ def call_method(
                 key_str = hash_key_vt.vt.as_python_constant()
                 tf_kwargs[key_str] = value_vt
 
-            output_old = tx.output.torch_function_enabled
             tx_old = tx.symbolic_torch_function_state.torch_function_subclass_enabled
-            tx.output.torch_function_enabled = False
             tx.symbolic_torch_function_state.torch_function_subclass_enabled = False
             try:
                 return func.call_function(tx, tf_args, tf_kwargs)
             finally:
-                tx.output.torch_function_enabled = output_old
                 tx.symbolic_torch_function_state.torch_function_subclass_enabled = (
                     tx_old
                 )
diff --git a/torch/_dynamo/variables/torch_function.py b/torch/_dynamo/variables/torch_function.py
@@ -629,7 +629,7 @@ def var_getattr(self, tx: "InstructionTranslator", name):
         # Handle non-overriden attributes inherited from `torch.Tensor`.
         attr_is_overriden = _is_attr_overidden(tx, self, name)
         if hasattr(torch.Tensor, name) and not attr_is_overriden:
-            if tx.output.torch_function_enabled:
+            if tx.symbolic_torch_function_state.torch_function_subclass_enabled:
                 if self.source:
                     install_guard(
                         AttrSource(
@@ -705,7 +705,7 @@ def call_method(
     ) -> "VariableTracker":
         # This code block implements inlining the __torch_function__ override
         # of `call_method`.
-        if tx.output.torch_function_enabled:
+        if tx.symbolic_torch_function_state.torch_function_subclass_enabled:
             import torch
 
             if _is_attr_overidden(tx, self, name):
diff --git a/torch/_guards.py b/torch/_guards.py
@@ -561,7 +561,6 @@ class GlobalContext(Checkpointable[GlobalContextCheckpointState]):
 
     _supported_global_states = {
         "grad_enabled",
-        "torch_function_enabled",
         "autocast_enabled",
         "autocast_cpu_enabled",
         "autocast_gpu_dtype",
