Revert "Change NamedTupleVariable implementation to subclass UserDefinedTupleVariable (pytorch#167468)"

This reverts commit c055ebe.

Reverted pytorch#167468 on behalf of https://github.com/pytorch-auto-revert due to Reverted automatically by pytorch's autorevert, to avoid this behaviour add the tag autorevert: disable ([comment](pytorch#167468 (comment)))

Author: pytorchmergebot

test/dynamo/test_functions.py

@@ -2053,7 +2053,7 @@ def test_namedtuple_defaults(a, b):
         return mytuple(tmp.x, tmp[1], tmp.xy + b)
 
     @make_test
-    def test_namedtuple_replace_1(a, b):
+    def test_namedtuple_replace(a, b):
         mytuple = collections.namedtuple("mytuple", ["x", "y"])
         t = mytuple(a, b)
         t._replace(x=b)
@@ -2109,7 +2109,7 @@ def test_namedtuple_user_methods(a, b):
         return mytuple.add(), mytuple.static_method(), mytuple.class_method()
 
     @make_test
-    def test_namedtuple_replace_2(a, b):
+    def test_namedtuple_replace(a, b):
         mytuple = FunctionTests.MyNamedTuple(a, b)
         replaced = mytuple._replace(first=b)
         return mytuple.first + mytuple.second + replaced.first + replaced.second

torch/_dynamo/variables/dicts.py

@@ -44,6 +44,7 @@
 )
 from .base import ValueMutationNew, VariableTracker
 from .constant import ConstantVariable
+from .lists import ListIteratorVariable
 
 
 if TYPE_CHECKING:
@@ -791,8 +792,6 @@ def call_method(
             self.call_method(tx, "update", args, kwargs)
             return self
         elif name == "__iter__":
-            from .lists import ListIteratorVariable
-
             if self.source and not is_constant_source(self.source):
                 tx.output.guard_on_key_order.add(self.source)
             return ListIteratorVariable(
@@ -1463,8 +1462,6 @@ def call_method(
         if name == "__len__":
             return self.dv_dict.call_method(tx, name, args, kwargs)
         elif name == "__iter__":
-            from .lists import ListIteratorVariable
-
             return ListIteratorVariable(
                 self.view_items_vt, mutation_type=ValueMutationNew()
             )

torch/_dynamo/variables/lists.py

@@ -15,6 +15,7 @@ class that handles its unique behaviors while integrating with Dynamo's
 """
 
 import collections
+import inspect
 import operator
 import sys
 from collections.abc import Sequence
@@ -38,6 +39,7 @@ class that handles its unique behaviors while integrating with Dynamo's
     get_fake_value,
     guard_if_dyn,
     iter_contains,
+    Lit,
     namedtuple_fields,
     odict_values,
     raise_args_mismatch,
@@ -46,8 +48,8 @@ class that handles its unique behaviors while integrating with Dynamo's
 )
 from .base import ValueMutationNew, VariableTracker
 from .constant import ConstantVariable
+from .functions import UserFunctionVariable, UserMethodVariable
 from .iter import IteratorVariable
-from .user_defined import UserDefinedTupleVariable
 
 
 if TYPE_CHECKING:
@@ -1294,51 +1296,24 @@ def call_obj_hasattr(
         return variables.ConstantVariable.create(hasattr(torch.Size, name))
 
 
-class NamedTupleVariable(UserDefinedTupleVariable):
+class NamedTupleVariable(TupleVariable):
     _nonvar_fields = {
         "tuple_cls",
         "dynamic_attributes",
-        *UserDefinedTupleVariable._nonvar_fields,
+        *TupleVariable._nonvar_fields,
     }
 
     def __init__(
         self,
         items: list[VariableTracker],
-        tuple_cls: type[tuple],
+        tuple_cls: type,
         dynamic_attributes: Optional[dict[str, VariableTracker]] = None,
         **kwargs: Any,
     ) -> None:
-        tuple_vt = variables.TupleVariable(
-            items, mutation_type=kwargs.get("mutation_type", ValueMutationNew())
-        )
-
-        # Create a dummy instance for method resolution
-        # This allows _maybe_get_baseclass_method to work correctly
-        fields = namedtuple_fields(tuple_cls)
-        num_fields = len(fields)
-        if tuple_cls.__module__ == "torch.return_types":
-            # Structseq: single iterable argument
-            dummy_value = tuple_cls([None] * num_fields)
-        else:
-            # Namedtuple: positional arguments
-            dummy_value = tuple_cls(*([None] * num_fields))  # type: ignore[arg-type]
-
-        super().__init__(
-            value=dummy_value,
-            tuple_vt=tuple_vt,
-            init_args=None,
-            **kwargs,
-        )
-
+        super().__init__(items, **kwargs)
         self.tuple_cls = tuple_cls
-        if len(self.tuple_cls.__mro__) < 3:
-            raise ValueError("NamedTuple should inherit from Tuple and Object.")
         self.dynamic_attributes = dynamic_attributes if dynamic_attributes else {}
 
-    @property
-    def items(self) -> list[VariableTracker]:
-        return self._tuple_vt.items
-
     def is_namedtuple(self) -> bool:
         return isinstance(getattr(self.tuple_cls, "_fields", None), tuple) and callable(
             getattr(self.tuple_cls, "_make", None)
@@ -1350,7 +1325,17 @@ def is_structseq(self) -> bool:
     def fields(self) -> tuple[str, ...]:
         return namedtuple_fields(self.tuple_cls)
 
-    def as_python_constant(self):
+    def debug_repr(self) -> str:
+        if self.is_structseq():
+            # StructSequenceType(iterable)
+            return repr(self.tuple_cls([Lit(x.debug_repr()) for x in self.items]))
+        # NamedTupleType(*iterable)
+        return repr(self.tuple_cls(*(Lit(x.debug_repr()) for x in self.items)))
+
+    def python_type(self) -> type:
+        return self.tuple_cls
+
+    def as_python_constant(self) -> Any:
         if self.is_structseq():
             # StructSequenceType(iterable)
             result = self.python_type()([x.as_python_constant() for x in self.items])
@@ -1372,39 +1357,57 @@ def as_python_constant(self):
 
         return result
 
-    def as_proxy(self):
+    def as_proxy(self) -> Any:
+        assert self.python_type() is not SizeVariable
         if self.is_structseq():
-            return self.python_type()([x.as_proxy() for x in self._tuple_vt.items])
-        return self.python_type()(*[x.as_proxy() for x in self._tuple_vt.items])
+            # StructSequenceType(iterable)
+            return self.python_type()(self._as_proxy())
+        # NamedTupleType(*iterable)
+        return self.python_type()(*self._as_proxy())
 
     def reconstruct(self, codegen: "PyCodegen") -> None:
+        # Always reconstruct the NamedTuple normally first
+        # Constructors:
+        #   StructSequenceType(iterable)
+        #   NamedTupleType(*iterable)
+        #   NamedTupleType._make(iterable)
         if self.is_structseq():
             create_fn = self.tuple_cls
         else:
             create_fn = self.tuple_cls._make  # type: ignore[attr-defined]
-
         codegen.add_push_null(
             lambda: codegen.append_output(
                 codegen.create_load_const_unchecked(create_fn)
             )
         )
-        codegen.foreach(self._tuple_vt.items)
+        codegen.foreach(self.items)
         codegen.extend_output(
             [
-                create_build_tuple(len(self._tuple_vt.items)),
+                create_build_tuple(len(self.items)),
             ]
             + create_call_function(1, False)
         )
 
-        # Apply initial dynamic attributes after construction (if any)
-        # Runtime dynamic attributes are tracked via side effects system
         for name, value in self.dynamic_attributes.items():
             codegen.dup_top()
             codegen(value)
             codegen.extend_output(create_rot_n(2))
             codegen.store_attr(name)
 
     def _is_method_overridden(self, method_name: str) -> bool:
+        """Checks if a method is overridden in the NamedTuple subclass.
+
+        Args:
+            method_name (str): The name of the method to check.
+
+        Returns:
+            bool: True if the method is overridden in the subclass, False otherwise.
+
+        Raises:
+            ValueError: If the NamedTuple class does not inherit from both Tuple and Object.
+        """
+        if len(self.tuple_cls.__mro__) < 3:
+            raise ValueError("NamedTuple should inherit from Tuple and Object.")
         if getattr(self.tuple_cls, method_name, None) == getattr(
             self.tuple_cls.__mro__[-3], method_name, None
         ):
@@ -1418,53 +1421,129 @@ def call_method(
         args: list[VariableTracker],
         kwargs: dict[str, VariableTracker],
     ) -> VariableTracker:
-        if self._is_method_overridden(name):
-            # Fall back to UserDefinedTupleVariable
-            return super().call_method(tx, name, args, kwargs)
-        elif name == "__setattr__":
+        if name == "__setattr__":
             if kwargs or len(args) != 2:
                 raise_args_mismatch(
                     tx,
                     name,
                     "2 args and 0 kwargs",
                     f"{len(args)} args and {len(kwargs)} kwargs",
                 )
-            attr_var, value = args
-            attr = attr_var.as_python_constant()
-
+            attr, value = args
+            attr = attr.as_python_constant()
             if (
                 # structseq is immutable
                 self.is_structseq()
                 # namedtuple directly created by `collections.namedtuple` is immutable
                 or self.tuple_cls.__bases__ == (tuple,)
+                # fields are immutable
                 or attr in self.fields()
             ):
                 raise_observed_exception(AttributeError, tx)
-
-            result = self.method_setattr_standard(tx, attr_var, value)
-            # Also update self.dynamic_attributes
+            # Subclass of namedtuple type can have dynamic attributes
+            tx.output.side_effects.mutation(self)
+            if self.source:
+                tx.output.side_effects.store_attr(self, attr, value)
             self.dynamic_attributes[attr] = value
-            return result
+            return ConstantVariable.create(None)
+        elif name == "_replace":
+            # NamedTuple._replace should create a new instance with replaced fields
+            if args:
+                raise_args_mismatch(tx, name, "0 args", f"{len(args)} args")
+
+            # Get the field names for validation
+            fields = self.fields()
+
+            # Start with current items (copy them)
+            new_items = list(self.items)
+
+            # Replace fields specified in kwargs
+            for field_name, new_value in kwargs.items():
+                if field_name not in fields:
+                    raise_observed_exception(
+                        ValueError,
+                        tx,
+                        args=[
+                            ConstantVariable.create(
+                                f"Got unexpected field name: '{field_name}'"
+                            )
+                        ],
+                    )
+
+                # Replace the item at the field's index
+                field_index = fields.index(field_name)
+                new_items[field_index] = new_value
+
+            return NamedTupleVariable(new_items, self.tuple_cls)
 
         return super().call_method(tx, name, args, kwargs)
 
-    def python_type(self) -> type:
-        return self.tuple_cls
+    def getitem_const(
+        self, tx: "InstructionTranslator", arg: VariableTracker
+    ) -> VariableTracker:
+        if isinstance(arg, SliceVariable):
+            # slicing a namedtuple produces a tuple
+            return TupleVariable(
+                self.items[arg.as_python_constant()],
+                source=None,
+            )
+        return super().getitem_const(tx, arg)
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        def check_and_create_method() -> Optional[VariableTracker]:
+            method = inspect.getattr_static(self.tuple_cls, name, None)
+            if isinstance(method, classmethod):
+                # We need the unbounded cls method to avoid the inline __self__
+                return UserMethodVariable(
+                    method.__func__,
+                    variables.UserDefinedClassVariable(self.tuple_cls),
+                )
+            elif isinstance(method, staticmethod):
+                # pyrefly: ignore[bad-argument-type]
+                return UserFunctionVariable(method.__func__)
+            elif inspect.isfunction(method):
+                return UserMethodVariable(method, self)
+            else:
+                return None
+
+        # Avoid UserMethodVariable fallback precisely when methods NamedTuple methods have not been overwritten.
+        if (
+            name == "_replace"
+            and not self._is_method_overridden("_replace")
+            and not self._is_method_overridden("__getattr__")
+        ):
+            # Return a BuiltinVariable for the _replace method
+            # Get the actual _replace method from the tuple class
+            actual_replace_method = getattr(self.tuple_cls, "_replace", None)
+            if actual_replace_method:
+                from ..source import AttrSource
+
+                source = AttrSource(self.source, name) if self.source else None
+                return variables.GetAttrVariable(self, name, source=source)
+            # Fallback if _replace doesn't exist (shouldn't happen for proper NamedTuples)
+            return super().var_getattr(tx, name)
 
-    def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
         if name == "_fields":
-            source = NamedTupleFieldsSource(self.source) if self.source else None
-            return VariableTracker.build(tx, self.fields(), source=source)
+            result_source = NamedTupleFieldsSource(self.source) if self.source else None
+            return VariableTracker.build(tx, self.fields(), source=result_source)
 
         if name in self.dynamic_attributes:
             return self.dynamic_attributes[name]
 
         fields = self.fields()
-        if name in fields:
-            field_index = fields.index(name)
-            return self._tuple_vt.items[field_index]
+        if name not in fields:
+            method = check_and_create_method()
+            if not method:
+                return super().var_getattr(tx, name)
+            return method
+        return self.items[fields.index(name)]
 
-        return super().var_getattr(tx, name)
+    def call_obj_hasattr(
+        self, tx: "InstructionTranslator", name: str
+    ) -> VariableTracker:
+        return variables.ConstantVariable.create(
+            name in self.dynamic_attributes or hasattr(self.tuple_cls, name)
+        )
 
 
 class SliceVariable(VariableTracker):