-
Notifications
You must be signed in to change notification settings - Fork 24.7k
[pytree] add another simplified pytree module torch.pytree
#148180
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: gh/XuehaiPan/249/base
Are you sure you want to change the base?
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148180
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Cancelled JobAs of commit 3e66bd8 with merge base bc6e066 ( CANCELLED JOB - The following job was cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: a073729 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: a073729 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: e6963b3 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: e6963b3 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 0570acc Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 2bdd0e3 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 0bf5096 Pull Request resolved: #148180
] | ||
|
||
|
||
def unflatten(treespec: PyTreeSpec, leaves: Iterable[_Any]) -> PyTree: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Use a TypeVar to dispatch iterable typing to the stub. No reason to erase the typing info here in case we improve typing in the future.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PyTree
is an alias of typing.Any
which is not a generic type.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
that might change in the futre though. Good to not erase typing if unnecessary and make it flexible to future refactors
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The change is localized and well documented + motivated, I'm happy with it
Thanks @XuehaiPan
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 9bac8a3 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: a88a117 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 4fd9651 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 0f7a5e7 Pull Request resolved: #148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: bab8a42 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 91cd5c0 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: bc6b68f Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: c42a1e5 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 024a983 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 024a983 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 024a983 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 024a983 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 704b18b Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 70a1538 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 70a1538 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 70a1538 Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 7c73c7a Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: 7c73c7a Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: e64ef6f Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: ac60aaa Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: ac60aaa Pull Request resolved: pytorch#148180
Differences between `torch.pytree` and `torch.utils.pytree`: 1. APIs in `torch.utils.pytree` have a `tree_` prefix: ```python leaves, treespec = torch.utils.pytree.tree_flatten(tree) new_tree = torch.utils.pytree.tree_map(func, tree) leaevs, treespec = torch.pytree.flatten(tree) new_tree = torch.pytree.map(func, tree) ``` 2. The argument order of `unflatten` is reversed for better `functools.partial` support: ```python tree = torch.utils.pytree.tree_unflatten(leaves, treespec) tree = torch.pytree.unflatten(treespec, leaves) unflatten_fn = functools.partial(torch.pytree.unflatten, treespec) tree1 = unflatten_fn(leaves1) tree2 = unflatten_fn(leaves2) ``` This is also aligned with the JAX pytree API: `jax.tree.unflatten(treedef, leaves)`. ghstack-source-id: acb3131 Pull Request resolved: pytorch#148180
Stack from ghstack (oldest at bottom):
__module__
#148328torch.pytree
#148180torch.utils.pytree
#137400tree_*
functions accept both Python and C++PyTreeSpec
#152624Differences between
torch.pytree
andtorch.utils.pytree
:APIs in
torch.utils.pytree
have atree_
prefix:This is similar to the JAX pytree API:
jax.tree_util.tree_*
vs.jax.tree.*
.The argument order of
unflatten
is reversed for betterfunctools.partial
support:This is also aligned with the JAX pytree API:
jax.tree.unflatten(treedef, leaves)
.Because we are adding a completely new module, there are no BC issues.
cc @zou3519