[dynamo][torch-function] Prevent unnecessary __torch_function__ tracing (#153551)

anijain2305 · pytorchmergebot · commit fa8543454ab5 · 2025-05-15T14:06:17.000Z
Pull Request resolved: #153551 Approved by: https://github.com/mlazos
diff --git a/test/dynamo_expected_failures/TestGradNewOnesOverride.test_newones b/test/dynamo_expected_failures/TestGradNewOnesOverride.test_newones
diff --git a/torch/_dynamo/polyfills/tensor.py b/torch/_dynamo/polyfills/tensor.py
@@ -11,25 +11,28 @@
 def make_subclass(
     cls: type[Any], data: torch.Tensor, requires_grad: bool = False, **kwargs: Any
 ) -> Any:
-    # This is a rough approximation of `THPVariable_make_subclass`. It should
-    # suffice for most of Dynamo tracing purposes.
-    # https://github.com/pytorch/pytorch/blob/ccfde4dadfa3c342076a1ee387017f84dd4ad2f7/torch/csrc/autograd/python_variable.cpp#L597-L650
-    assert len(kwargs) == 0, "_make_subclass only supports requires_grad as keyword arg"
-    data = data.detach()
-
-    # Avoid unnecessary `requires_grad` mutation, which isn't supported in Dynamo.
-    if data.requires_grad != requires_grad:
-        data.requires_grad = requires_grad
-
-    # Dynamo can't yet handle upcasting to base tensor type via `as_subclass`.
-    if cls is torch.Tensor:
-        return torch.Tensor(data)
-
-    # Calling `as_subclass` because
-    # 1. Dynamo knows how to handle it
-    # 2. the C impls match at this point -- both `THPVariable_make_subclass` and
-    #    `THPVariable_as_subclass` calls `THPVariable_NewWithVar`.
-    return data.as_subclass(cls)
+    with torch._C.DisableTorchFunctionSubclass():
+        # This is a rough approximation of `THPVariable_make_subclass`. It should
+        # suffice for most of Dynamo tracing purposes.
+        # https://github.com/pytorch/pytorch/blob/ccfde4dadfa3c342076a1ee387017f84dd4ad2f7/torch/csrc/autograd/python_variable.cpp#L597-L650
+        assert len(kwargs) == 0, (
+            "_make_subclass only supports requires_grad as keyword arg"
+        )
+        data = data.detach()
+
+        # Avoid unnecessary `requires_grad` mutation, which isn't supported in Dynamo.
+        if data.requires_grad != requires_grad:
+            data.requires_grad = requires_grad
+
+        # Dynamo can't yet handle upcasting to base tensor type via `as_subclass`.
+        if cls is torch.Tensor:
+            return torch.Tensor(data)
+
+        # Calling `as_subclass` because
+        # 1. Dynamo knows how to handle it
+        # 2. the C impls match at this point -- both `THPVariable_make_subclass` and
+        #    `THPVariable_as_subclass` calls `THPVariable_NewWithVar`.
+        return data.as_subclass(cls)
 
 
 __all__ = [
diff --git a/torch/_dynamo/variables/torch_function.py b/torch/_dynamo/variables/torch_function.py
@@ -642,7 +642,11 @@ 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 (
+            hasattr(torch.Tensor, name)
+            and not attr_is_overriden
+            and not inspect.ismethoddescriptor(getattr(torch.Tensor, name))
+        ):
             args, kwargs = [self], {}
             if can_dispatch_torch_function(tx, args, kwargs):
                 if self.source:
