Properly detect nested torch function args (#127496)

Summary:
Dynamo was not detecting nested torch function classes in containers. This was due to pytree compatibility for variable trackers being removed.
Fixes pytorch/pytorch#127174

X-link: pytorch/pytorch#127496
Approved by: https://github.com/anijain2305

Reviewed By: atalman

Differential Revision: D58110097

Pulled By: mlazos

fbshipit-source-id: fb831e7237fb0d115a435571539e4eae3dad2d52

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

@@ -2529,12 +2529,17 @@ def is_torch_function_object(value):
 
 
 def has_torch_function(vt: "torch._dynamo.variables.base.VariableTracker") -> bool:
- from torch._dynamo.variables import UserDefinedObjectVariable
+ from torch._dynamo.variables import LazyVariableTracker, UserDefinedObjectVariable
  from torch._dynamo.variables.torch_function import TensorWithTFOverrideVariable
 
- return isinstance(vt, TensorWithTFOverrideVariable) or (
- isinstance(vt, UserDefinedObjectVariable)
- and hasattr(vt.value, "__torch_function__")
+ if isinstance(vt, TensorWithTFOverrideVariable):
+ return True
+
+ if isinstance(vt, LazyVariableTracker):
+ LazyVariableTracker.realize(vt)
+
+ return isinstance(vt, UserDefinedObjectVariable) and hasattr(
+ vt.value, "__torch_function__"
  )