[dynamo][user-defined] Unify standard and non-standard __new__ codebase

anijain2305 · anijain2305 · commit af0f5834bac8 · 2025-02-09T00:32:51.000-08:00
ghstack-source-id: 6338f8d Pull Request resolved: #146737
diff --git a/torch/_dynamo/variables/user_defined.py b/torch/_dynamo/variables/user_defined.py
@@ -607,17 +607,6 @@ def call_function(
             )
             var.call_method(tx, "__init__", args, kwargs)
             return var
-        elif (
-            self.is_standard_new()
-            and SideEffects.cls_supports_mutation_side_effects(self.value)
-            and self.source
-        ):
-            var = tx.output.side_effects.track_new_user_defined_object(
-                variables.BuiltinVariable(object), self, args
-            )
-            with do_not_convert_to_tracable_parameter():
-                var.call_method(tx, "__init__", args, kwargs)
-                return var
         elif (
             variables.RestrictedListSubclassVariable.is_matching_cls(self.value)
             and self.source
@@ -678,18 +667,15 @@ def call_function(
             # types.MappingProxyType is a read-only proxy of the dict. If the
             # original dict changes, the changes are reflected in proxy as well.
             return variables.MappingProxyVariable(args[0])
-        elif (
-            not self.is_standard_new()
-            and SideEffects.cls_supports_mutation_side_effects(self.value)
-            and self.source
-        ):
-            return tx.inline_user_function_return(
-                VariableTracker.build(
-                    tx, polyfills.instantiate_user_defined_class_object
-                ),
-                [self, *args],
-                kwargs,
-            )
+        elif SideEffects.cls_supports_mutation_side_effects(self.value) and self.source:
+            with do_not_convert_to_tracable_parameter():
+                return tx.inline_user_function_return(
+                    VariableTracker.build(
+                        tx, polyfills.instantiate_user_defined_class_object
+                    ),
+                    [self, *args],
+                    kwargs,
+                )
         return super().call_function(tx, args, kwargs)
 
     def is_standard_new(self):
@@ -1179,6 +1165,11 @@ def var_getattr(self, tx: "InstructionTranslator", name):
             elif getattr_fn is not None:
                 unimplemented("UserDefined with non-function __getattr__")
 
+        from ..mutation_guard import unpatched_nn_module_init
+
+        if subobj is torch.nn.Module.__init__:
+            subobj = unpatched_nn_module_init
+
         if isinstance(subobj, property):
             if self.source:
                 # Read the class attribute to reach the property
