Update base for Update on "[dynamo][not ready] polyfill infra for classes"

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]
This commit is contained in:
Animesh Jain 2025-02-09 17:56:29 -08:00
parent 0da1d3ce27
commit b80a86ba0f

View file

@ -297,6 +297,7 @@ class SideEffects:
from .variables.ctx_manager import GenericContextWrappingVariable
from .variables.torch_function import TorchFunctionModeVariable
from .variables.user_defined import is_forbidden_context_manager
variable_cls: type[
variables.UserDefinedObjectVariable
@ -305,9 +306,13 @@ class SideEffects:
user_cls, TorchFunctionMode
) and TorchFunctionModeVariable.is_supported_torch_function_mode(user_cls):
variable_cls = TorchFunctionModeVariable
elif hasattr(user_cls, "__enter__") and hasattr(user_cls, "__exit__"):
elif (
hasattr(user_cls, "__enter__")
and hasattr(user_cls, "__exit__")
and not is_forbidden_context_manager(user_cls)
):
variable_cls = GenericContextWrappingVariable
if issubclass(user_cls, torch.nn.Module):
elif issubclass(user_cls, torch.nn.Module):
variable_cls = variables.UnspecializedNNModuleVariable
elif issubclass(user_cls, (dict, collections.OrderedDict)):
variable_cls = variables.UserDefinedDictVariable