[pytree] add another simplified pytree module `torch.pytree` #148180

XuehaiPan · 2025-02-28T10:51:49Z

Stack from ghstack (oldest at bottom):

Differences between torch.pytree and torch.utils.pytree:

APIs in torch.utils.pytree have a tree_ prefix:

leaves, treespec = torch.utils.pytree.tree_flatten(tree)
new_tree = torch.utils.pytree.tree_map(func, tree)

leaves, treespec = torch.pytree.flatten(tree)
new_tree = torch.pytree.map(func, tree)

This is similar to the JAX pytree API: jax.tree_util.tree_* vs. jax.tree.*.

The argument order of unflatten is reversed for better functools.partial support:

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).

Because we are adding a completely new module, there are no BC issues.

cc @zou3519

[ghstack-poisoned]

pytorch-bot · 2025-02-28T10:51:53Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148180

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

CI workflows being skipped on PR

✅ No Failures

As of commit 41259dc with merge base e9e1aac ():
💚 Looks good so far! There are no failures yet. 💚

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

[ghstack-poisoned]

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

[ghstack-poisoned]

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

[ghstack-poisoned]

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

[ghstack-poisoned]

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

Skylion007 · 2025-02-28T14:54:06Z

torch/pytree.py

+]
+
+
+def unflatten(treespec: PyTreeSpec, leaves: Iterable[_Any]) -> PyTree:


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.

PyTree is an alias of typing.Any which is not a generic type.

that might change in the futre though. Good to not erase typing if unnecessary and make it flexible to future refactors

vmoens

The change is localized and well documented + motivated, I'm happy with it
Thanks @XuehaiPan

torch/pytree.py

[ghstack-poisoned]

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

[ghstack-poisoned]

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

[ghstack-poisoned]

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

[ghstack-poisoned]

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

[ghstack-poisoned]

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: 7f55994 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 511d44e Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: d4b7c9b Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 5349d33 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 15e86f3 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: f9e47cf Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 5b8af67 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: be02f1a Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 2185a81 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 6e3fce8 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: 3d16082 Pull Request resolved: pytorch#148180

[ghstack-poisoned]

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: f1bb7f0 Pull Request resolved: pytorch#148180

Update

3ae6494

[ghstack-poisoned]

XuehaiPan mentioned this pull request Feb 28, 2025

[pytree] Add public pytree module torch.utils.pytree #137400

Open

XuehaiPan added topic: not user facing topic category module: pytree labels Feb 28, 2025

XuehaiPan requested a review from zou3519 February 28, 2025 10:55

pytorchbot added the open source label Feb 28, 2025

Update

51232f8

[ghstack-poisoned]

Update

c5f436b

[ghstack-poisoned]

Update

c36201b

[ghstack-poisoned]

XuehaiPan requested a review from albanD as a code owner February 28, 2025 14:11

Update

08affe6

[ghstack-poisoned]

Skylion007 reviewed Feb 28, 2025

View reviewed changes

zou3519 requested a review from vmoens February 28, 2025 15:19

vmoens approved these changes Feb 28, 2025

View reviewed changes

torch/pytree.py Outdated Show resolved Hide resolved

torch/pytree.py Outdated Show resolved Hide resolved

Update

4e773f0

[ghstack-poisoned]

Update

308f6a0

[ghstack-poisoned]

Update

00a6482

[ghstack-poisoned]

Update

f62063d

[ghstack-poisoned]

Update

c960761

[ghstack-poisoned]

XuehaiPan added 2 commits April 15, 2025 22:10

Update

97cf576

[ghstack-poisoned]

Update

898c203

[ghstack-poisoned]

Update

0d8d2a0

[ghstack-poisoned]

Update

ad0668c

[ghstack-poisoned]

Update

bfe4839

[ghstack-poisoned]

XuehaiPan added 2 commits May 2, 2025 01:39

Update

c309091

[ghstack-poisoned]

Update

fffb30d

[ghstack-poisoned]

Update

599b396

[ghstack-poisoned]

XuehaiPan mentioned this pull request May 1, 2025

[pytree] make tree_* functions accept both Python and C++ PyTreeSpec #152624

Open

Update

5710b9d

[ghstack-poisoned]

XuehaiPan added 2 commits May 3, 2025 00:40

Update

313bfee

[ghstack-poisoned]

Update

ba32ef9

[ghstack-poisoned]

Update

dbd47d2

[ghstack-poisoned]

Update

113b830

[ghstack-poisoned]

Update

41259dc

[ghstack-poisoned]

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[pytree] add another simplified pytree module `torch.pytree` #148180

[pytree] add another simplified pytree module `torch.pytree` #148180

XuehaiPan commented Feb 28, 2025 •

edited

Loading

pytorch-bot bot commented Feb 28, 2025 •

edited

Loading

Skylion007 Feb 28, 2025

XuehaiPan Feb 28, 2025

Skylion007 Feb 28, 2025

vmoens left a comment

		]


		def unflatten(treespec: PyTreeSpec, leaves: Iterable[_Any]) -> PyTree:

[pytree] add another simplified pytree module torch.pytree #148180

Are you sure you want to change the base?

[pytree] add another simplified pytree module torch.pytree #148180

Conversation

XuehaiPan commented Feb 28, 2025 • edited Loading

pytorch-bot bot commented Feb 28, 2025 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148180

❗ 1 Active SEVs

✅ No Failures

Skylion007 Feb 28, 2025

Choose a reason for hiding this comment

XuehaiPan Feb 28, 2025

Choose a reason for hiding this comment

Skylion007 Feb 28, 2025

Choose a reason for hiding this comment

vmoens left a comment

Choose a reason for hiding this comment

[pytree] add another simplified pytree module `torch.pytree` #148180

[pytree] add another simplified pytree module `torch.pytree` #148180

XuehaiPan commented Feb 28, 2025 •

edited

Loading

pytorch-bot bot commented Feb 28, 2025 •

edited

Loading