Update on "[pytree] add key path api"

suo · suo · commit f642001e7e20 · 2024-01-16T12:11:06.000-08:00
This PR introduces a key path API to pytrees, drawing direct inspiration from JAX's [key path API](https://jax.readthedocs.io/en/latest/jax-101/05.1-pytrees.html#key-paths). I added the 3 APIs described there, and a registry of `flatten_with_keys` fns for each node type, which is a version of `flatten` that also returns `KeyEntry`s describing how to access values from the original pytree. Current use cases for this API: - Folks would like to do argument traversal over input pytrees to do verification and compatibility enforcement. Keypaths are useful for this—https://fburl.com/code/06p7zrvr is a handrolled pass doing basically the same thing but probably more fragilely. - In export non-strict mode, we need to figure out a way to track sources for pytree inputs. In strict mode, dynamo handles this for us, but we'd like a decoupled component to handle this when we're not using dynamo. I'm sure there are places it would be useful. Some design notes: - I only implemented the API for the Python pytree impl. optree has some differences in how their keypath APIs are designed (see #113378 for discussion). I have some issues with the proposed typed_path solution in that discussion and prefer JAX's API, but we can hash that out separately. - The way folks register a `flatten_with_keys` fn is through a new kwarg to `register_pytree_node`. This follows how we do serialization fns, although the list of additional arguments is getting unwieldy. - My impl handles pytrees with an undefined `flatten_with_keys` fn is different from JAX. I will raise an error, JAX creates a fallback keyentry. Differential Revision: [D52547850](https://our.internmc.facebook.com/intern/diff/D52547850/) [ghstack-poisoned]
diff --git a/test/test_pytree.py b/test/test_pytree.py
@@ -1134,6 +1134,36 @@ class ACustomPytree:
             actual = py_pytree.tree_unflatten([leaf for _, leaf in key_leaves], spec)
             self.assertEqual(actual, pytree)
 
+    def test_tree_leaves_with_path(self):
+        class ANamedTuple(NamedTuple):
+            x: torch.Tensor
+            y: int
+            z: str
+
+        @dataclass
+        class ACustomPytree:
+            x: Any
+            y: Any
+            z: Any
+
+        py_pytree.register_pytree_node(
+            ACustomPytree,
+            flatten_fn=lambda f: ([f.x, f.y], f.z),
+            unflatten_fn=lambda xy, z: ACustomPytree(xy[0], xy[1], z),
+            flatten_with_keys_fn=lambda f: ((("x", f.x), ("y", f.y)), f.z),
+        )
+
+        SOME_PYTREES = [
+            (None,),
+            ["hello", [1, 2], {"foo": [(3)]}],
+            [ANamedTuple(x=torch.rand(2, 3), y=1, z="foo")],
+            [ACustomPytree(x=12, y={"cin": [1, 4, 10], "bar": 18}, z="leaf"), 5],
+        ]
+        for pytree in SOME_PYTREES:
+            flat_out, _ = py_pytree.tree_flatten_with_path(pytree)
+            leaves_out = py_pytree.tree_leaves_with_path(pytree)
+            self.assertEqual(flat_out, leaves_out)
+
     def test_key_str(self):
         class ANamedTuple(NamedTuple):
             x: str
