tree_flatten supports named tuples (#1446)

thunder/core/baseutils.py

@@ -210,6 +210,17 @@ def get_module(name: str) -> Any:
     return sys.modules[name]
 
 
+def is_likely_from_collections_namedtuple(tuple_type):
+    from collections import namedtuple
+
+    # Check if tuple_type code object is coming from namedtuple
+    return (
+        hasattr(tuple_type, "__repr__")
+        and hasattr(tuple_type.__repr__, "__code__")
+        and tuple_type.__repr__.__code__ in namedtuple.__code__.co_consts
+    )
+
+
 #
 # Functions related to printing and debugging
 #

thunder/core/interpreter.py

@@ -42,7 +42,7 @@
     TracebackType,
 )
 
-from thunder.core.baseutils import Singleton, init_colors, extract_callable_name
+from thunder.core.baseutils import Singleton, init_colors, extract_callable_name, is_likely_from_collections_namedtuple
 from thunder.core.codeutils import Positions
 
 
@@ -2848,16 +2848,6 @@ def _tuple_new_provenance_tracking_lookaside(cls, iterable=(), /):
             else:
                 item_wrappers.append(wv)
 
-    def is_likely_from_collections_namedtuple(tuple_type):
-        from collections import namedtuple
-
-        # Check if tuple_type code object is coming from namedtuple
-        return (
-            hasattr(tuple_type, "__repr__")
-            and hasattr(tuple_type.__repr__, "__code__")
-            and tuple_type.__repr__.__code__ in namedtuple.__code__.co_consts
-        )
-
     # Construction of namedtuples may raise
     try:
         ures = tuple(w.value for w in item_wrappers)

thunder/core/pytree.py

@@ -6,7 +6,7 @@
 import torch
 import thunder.core.dtypes as dtypes
 import thunder.core.devices as devices
-from thunder.core.baseutils import ProxyInterface
+from thunder.core.baseutils import ProxyInterface, is_likely_from_collections_namedtuple
 from types import FunctionType
 
 OPTREE_NAMESPACE = "thunder"
@@ -61,6 +61,7 @@ def tree_flatten(args, namespace=OPTREE_NAMESPACE):
             torch.autograd.function.FunctionCtx,
         }
         and not isinstance(args, (ProxyInterface))
+        and not is_likely_from_collections_namedtuple(args)
         and not dataclasses.is_dataclass(args)
         and not type(args).__module__.startswith("torch.return_types")
     ):

thunder/tests/test_core.py

@@ -662,6 +662,17 @@ def test_to_printable_not_collection():
         assert inp is out
 
 
+def test_to_printable_collection():
+    from collections import namedtuple
+
+    MyTuple = namedtuple("MyTuple", ["x", "y"])
+
+    inps = (MyTuple("abc", "def"),)
+    for inp in inps:
+        out = codeutils.to_printable(None, inp)
+        assert inp == out
+
+
 #
 # Type promotion tests
 #