Teach pytrees about namedtuple (pytorch#62292)

Summary:
Pull Request resolved: pytorch#62292

This PR adds pytree support for namedtuples. The challenge about namedtuple
is that each namedtuple class is actually different. This PR does the
following:
- it adds a namedtuple flatten/unflatten. The flatten function returns
a context that is the actual type of the namedtuple subclass. The
unflatten function uses that type to reconstruct the namedtuple
- Special cases all pytree logic to consider all namedtuples the same.
This is done by creating a `_get_node_type(pytree)` helper function that
returns `namedtuple` if `pytree` is any namedtuple subclass. The effect
of this is that all namedtuple subclasses will go through the namedtuple
flatten/unflatten functions
- Adds a `_namedtuple_flatten_spec` function for FX pytrees. This function
flattens the namedtuple based on the spec and is equivalent to the
`_tuple_flatten_spec`.

Test Plan
- new tests in test/test_pytree.py and test/test_fx.py

Test Plan: Imported from OSS

Reviewed By: albanD

Differential Revision: D29947302

Pulled By: zou3519

fbshipit-source-id: 19c00665b13546642c315df0f243ad99b8e7ff7c

Author: zou3519

test/test_fx.py

@@ -28,6 +28,7 @@
 from torch.fx.experimental.rewriter import RewritingTracer
 from torch.fx.operator_schemas import get_signature_for_torch_op
 from copy import deepcopy
+from collections import namedtuple
 
 from torch.fx.proxy import TraceError
 
@@ -65,6 +66,11 @@ def forward(self, x):
 def a_non_torch_leaf(a, b):
     return a + b
 
+# used in test_pytree. It's all the way out here because pickling a GraphModule
+# that uses Point errors out if Point is local to the function
+Point = namedtuple('Point', ['x', 'y'])
+
+
 # Test wrap() passing both a function name as well as a function
 # directly
 def a_lifted_leaf(a, b):
@@ -2610,6 +2616,8 @@ def f_dict_list_map(x):
         def f_dict_add(x):
             return x['a'] + sum(x['z'])
 
+        def f_namedtuple_add(x):
+            return x.x + x.y
 
         pytree._register_pytree_node(
             Foo,
@@ -2639,6 +2647,7 @@ def f_return_custom(x):
             (f_custom, Foo(PH, 3)),
             (f_custom_dict, Foo({'a': PH, 'b': PH}, PH)),
             # (f_return_custom, Foo(PH, PH)), # Don't currently support output pytrees
+            (f_namedtuple_add, Point(PH, PH)),
         ]
 
         def verify_pytree(f, inp):

test/test_pytree.py

@@ -2,6 +2,7 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten, TreeSpec, LeafSpec
 from torch.utils._pytree import _broadcast_to_and_flatten
+from collections import namedtuple
 
 class TestPytree(TestCase):
     def test_treespec_equality(self):
@@ -59,6 +60,33 @@ def run_test(tup):
         run_test((1., 2))
         run_test((torch.tensor([1., 2]), 2, 10, 9, 11))
 
+    def test_flatten_unflatten_namedtuple(self):
+        Point = namedtuple('Point', ['x', 'y'])
+
+        def run_test(tup):
+            expected_spec = TreeSpec(namedtuple, Point, [LeafSpec() for _ in tup])
+            values, treespec = tree_flatten(tup)
+            self.assertTrue(isinstance(values, list))
+            self.assertEqual(values, list(tup))
+            self.assertEqual(treespec, expected_spec)
+
+            unflattened = tree_unflatten(values, treespec)
+            self.assertEqual(unflattened, tup)
+            self.assertTrue(isinstance(unflattened, Point))
+
+        run_test(Point(1., 2))
+        run_test(Point(torch.tensor(1.), 2))
+
+    def test_flatten_unflatten_torch_namedtuple_return_type(self):
+        x = torch.randn(3, 3)
+        expected = torch.max(x, dim=0)
+
+        values, spec = tree_flatten(expected)
+        result = tree_unflatten(values, spec)
+
+        self.assertEqual(type(result), type(expected))
+        self.assertEqual(result, expected)
+
     def test_flatten_unflatten_dict(self):
         def run_test(tup):
             expected_spec = TreeSpec(dict, list(tup.keys()),

torch/fx/_pytree.py

@@ -1,5 +1,6 @@
-from typing import Callable, Any, Tuple, List, Dict, Type
+from typing import Callable, Any, Tuple, List, Dict, Type, NamedTuple
 from torch.utils._pytree import PyTree, TreeSpec, LeafSpec
+from collections import namedtuple
 
 FlattenFuncSpec = Callable[[PyTree, TreeSpec], List]
 
@@ -32,6 +33,10 @@ def _list_flatten_spec(d: List[Any], spec: TreeSpec) -> List[Any]:
 def _tuple_flatten_spec(d: Tuple[Any], spec: TreeSpec) -> List[Any]:
     return [d[i] for i in range(len(spec.children_specs))]
 
+def _namedtuple_flatten_spec(d: NamedTuple, spec: TreeSpec) -> List[Any]:
+    return [d[i] for i in range(len(spec.children_specs))]
+
 register_pytree_flatten_spec(dict, _dict_flatten_spec)
 register_pytree_flatten_spec(list, _list_flatten_spec)
 register_pytree_flatten_spec(tuple, _tuple_flatten_spec)
+register_pytree_flatten_spec(namedtuple, _tuple_flatten_spec)

torch/utils/_pytree.py

@@ -1,4 +1,5 @@
 from typing import NamedTuple, Callable, Any, Tuple, List, Dict, Type, cast, Optional
+from collections import namedtuple
 
 """
 Contains utility functions for working with nested python data structures.
@@ -56,14 +57,38 @@ def _tuple_flatten(d: Tuple[Any, ...]) -> Tuple[List[Any], Context]:
 def _tuple_unflatten(values: List[Any], context: Context) -> Tuple[Any, ...]:
     return tuple(values)
 
+def _namedtuple_flatten(d: NamedTuple) -> Tuple[List[Any], Context]:
+    return list(d), type(d)
+
+def _namedtuple_unflatten(values: List[Any], context: Context) -> NamedTuple:
+    return cast(NamedTuple, context(*values))
+
 _register_pytree_node(dict, _dict_flatten, _dict_unflatten)
 _register_pytree_node(list, _list_flatten, _list_unflatten)
 _register_pytree_node(tuple, _tuple_flatten, _tuple_unflatten)
+_register_pytree_node(namedtuple, _namedtuple_flatten, _namedtuple_unflatten)
+
 
+# h/t https://stackoverflow.com/questions/2166818/how-to-check-if-an-object-is-an-instance-of-a-namedtuple
+def _is_namedtuple_instance(pytree: Any) -> bool:
+    typ = type(pytree)
+    bases = typ.__bases__
+    if len(bases) != 1 or bases[0] != tuple:
+        return False
+    fields = getattr(typ, '_fields', None)
+    if not isinstance(fields, tuple):
+        return False
+    return all(type(entry) == str for entry in fields)
+
+def _get_node_type(pytree: Any) -> Any:
+    if _is_namedtuple_instance(pytree):
+        return namedtuple
+    return type(pytree)
 
 # A leaf is defined as anything that is not a Node.
 def _is_leaf(pytree: PyTree) -> bool:
-    return type(pytree) not in SUPPORTED_NODES.keys()
+    return _get_node_type(pytree) not in SUPPORTED_NODES.keys()
+
 
 # A TreeSpec represents the structure of a pytree. It holds:
 # "type": the type of root Node of the pytree
@@ -105,7 +130,8 @@ def tree_flatten(pytree: PyTree) -> Tuple[List[Any], TreeSpec]:
     if _is_leaf(pytree):
         return [pytree], LeafSpec()
 
-    flatten_fn = SUPPORTED_NODES[type(pytree)].flatten_fn
+    node_type = _get_node_type(pytree)
+    flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
     child_pytrees, context = flatten_fn(pytree)
 
     # Recursively flatten the children
@@ -116,7 +142,7 @@ def tree_flatten(pytree: PyTree) -> Tuple[List[Any], TreeSpec]:
         result += flat
         children_specs.append(child_spec)
 
-    return result, TreeSpec(type(pytree), context, children_specs)
+    return result, TreeSpec(node_type, context, children_specs)
 
 
 def tree_unflatten(values: List[Any], spec: TreeSpec) -> PyTree:
@@ -167,10 +193,11 @@ def _broadcast_to_and_flatten(pytree: PyTree, spec: TreeSpec) -> Optional[List[A
         return [pytree] * spec.num_leaves
     if isinstance(spec, LeafSpec):
         return None
-    if type(pytree) != spec.type:
+    node_type = _get_node_type(pytree)
+    if node_type != spec.type:
         return None
 
-    flatten_fn = SUPPORTED_NODES[type(pytree)].flatten_fn
+    flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
     child_pytrees, ctx = flatten_fn(pytree)
 
     # Check if the Node is different from the spec