8000 [pytree] Add public pytree module `torch.utils.pytree` by XuehaiPan · Pull Request #137400 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

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

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.

Sign up for GitHub

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

Jump to bottom
Open
wants to merge 158 commits into
base: gh/XuehaiPan/179/base
Choose a base branch
from
Open

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

wants to merge 158 commits into from

Conversation

XuehaiPan
Copy link
Collaborator
@XuehaiPan XuehaiPan commented Oct 5, 2024 •
edited
Loading

Stack from ghstack (oldest at bottom):

Populate APIs from torch.utils._pytree (default) or torch.utils._cxx_pytree to a new public module torch.utils.pytree. There is an environment variable PYTORCH_USE_CXX_PYTREE (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using torch.utils.pytree.

Resolves #65761

cc @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @rec

@XuehaiPan
Update
b69d677 
[ghstack-poisoned]
@pytorch-bot PyTorch Bot
Copy link
pytorch-bot bot commented Oct 5, 2024 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit 087c7b2 with merge base cb5f31a (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

This was referenced Oct 5, 2024
Closed
Closed
Closed
Closed
Closed
XuehaiPan added a commit that referenced this pull request Oct 5, 2024
@XuehaiPan
[pytree][4/N] make torch.utils.pytree as public API
a4ef04a 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: ea0d2b6
Pull Request resolved: #137400
@XuehaiPan XuehaiPan self-assigned this Oct 5, 2024
@XuehaiPan XuehaiPan added ciflow/trunk Trigger trunk jobs on your pull request release notes: releng release notes category module: pytree release notes: dynamo labels Oct 5, 2024
@XuehaiPan XuehaiPan linked an issue Oct 5, 2024 that may be closed by this pull request
torch.utils._pytree -> stable #65761
Open
@XuehaiPan
Update
b864295 
[ghstack-poisoned]
XuehaiPan added a commit that referenced this pull request Oct 5, 2024
@XuehaiPan
[pytree][4/N] make torch.utils.pytree as public API
0fb0566 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 656e242
Pull Request resolved: #137400
@XuehaiPan
Update
bdb2723 
[ghstack-poisoned]
XuehaiPan added a commit that referenced this pull request Oct 5, 2024
@XuehaiPan
[pytree][4/N] make torch.utils.pytree as public API
078eee4 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: db25160
Pull Request resolved: #137400
@XuehaiPan
Update
9b21fd9 
[ghstack-poisoned]
XuehaiPan added a commit that referenced this pull request Oct 5, 2024
@XuehaiPan
[pytree][4/N] make torch.utils.pytree as public API
cd84dfe 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 65bf5dc
Pull Request resolved: #137400
@XuehaiPan
Update
f1f9fe1 
[ghstack-poisoned]
XuehaiPan added a commit that referenced this pull request Oct 5, 2024
@XuehaiPan
[pytree][4/N] make torch.utils.pytree as public API
0c92aea 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 2718a77
Pull Request resolved: #137400
@XuehaiPan
Update
d98b761 
[ghstack-poisoned]
XuehaiPan added a commit that referenced this pull request Oct 5, 2024
@XuehaiPan
[pytree][4/N] make torch.utils.pytree as public API
d1450c7 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 57d7a7b
Pull Request resolved: #137400
@XuehaiPan
Update
4ae51d3 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Apr 26, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
12b15d3 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 000d799
Pull Request resolved: pytorch#137400
angelayi
angelayi reviewed May 1, 2025
View reviewed changes
torch/utils/pytree/__init__.py
__all__ = [
"PyTreeSpec",
"register_pytree_node",
"tree_flatten",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what about tree_flatten_with_spec and register_dataclass?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure if they are exposed APIs or only used internally. If it should be exposed and stabilized in torch.utils.pytree, I can add it in a follow-up.

@XuehaiPan
Update
f72b55e 
[ghstack-poisoned]
@XuehaiPan XuehaiPan requested a review from angelayi May 1, 2025 17:41
@XuehaiPan
Update
c91f575 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 1, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
7666458 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 42d47d5
Pull Request resolved: pytorch#137400
zou3519
zou3519 reviewed May 1, 2025
View reviewed changes
torch/utils/__init__.py
@@ -11,6 +11,7 @@
data as data,
deterministic as deterministic,
hooks as hooks,
pytree as pytree,
Copy link
Contributor
@zou3519 zou3519 May 1, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because we are offering an API to swap between python and C++ pytree in this PR, for this PR, I would like us to have finished making Python treespec and C++ treespec interoperable.

That is, one should be able to pass a treespec created from Python pytree to C++-pytree and vice versa. Otherwise there will be pain switching between the two APIs.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Prepended a new PR in the stack to resolve this:

@XuehaiPan
Update
7e0f546 
[ghstack-poisoned]
@XuehaiPan XuehaiPan mentioned this pull request May 1, 2025
Open
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 1, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
ab3c134 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 4778e60
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
abae549 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 1, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
2811bd9 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 684f8c3
Pull Request resolved: pytorch#137400
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 2, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
3987a1d 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 684f8c3
Pull Request resolved: pytorch#137400
XuehaiPan added 2 commits May 3, 2025 00:40
@XuehaiPan
Update
2fc047b 
[ghstack-poisoned]
@XuehaiPan
Update
e635263 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 2, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
bf8b0f9 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 5d2275d
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
9efa19c 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 2, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
802f305 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 952db83
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
81e6b01 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 2, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
4c3a694 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 7db73f6
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
d665e57 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 3, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
526dac5 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 661af0f
Pull Request resolved: pytorch#137400
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 5, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
a67ca72 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 661af0f
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
4b154fa 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 8, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
5dedd77 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 0ead633
Pull Request resolved: pytorch#137400
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 11, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
c9a8feb 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: 0ead633
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
a2bb393 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 14, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
e1643bc 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: e075efa
Pull Request resolved: pytorch#137400
@XuehaiPan
Update
087c7b2 
[ghstack-poisoned]
XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 16, 2025
@XuehaiPan
[pytree] make torch.utils.pytree as public API
3018dcc 
Populate APIs from `torch.utils._pytree` (default) or `torch.utils._cxx_pytree` to a new public module `torch.utils.pytree`. There is a environment varaible `PYTORCH_USE_CXX_PYTREE` (disabled by default) to control this.

Since the CXX pytree is now Dynamo traceable, the users can change the underlining pytree implementation by flipping the environment variable while using `torch.utils.pytree`.

ghstack-source-id: f21d8ee
Pull Request resolved: pytorch#137400
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@zou3519 zou3519 zou3519 left review comments

@malfet < AE87 a class="assignee Link--primary css-truncate-target width-fit" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="/malfet"> malfet Awaiting requested review from malfet

@ezyang ezyang Awaiting requested review from ezyang

@albanD albanD Awaiting requested review from albanD

@jeffdaily jeffdaily Awaiting requested review from jeffdaily

@avikchaudhuri avikchaudhuri Awaiting requested review from avikchaudhuri avikchaudhuri is a code owner

@tugsbayasgalan tugsbayasgalan Awaiting requested review from tugsbayasgalan tugsbayasgalan is a code owner

@zhxchen17 zhxchen17 Awaiting requested review from zhxchen17 zhxchen17 is a code owner

@ydwu4 ydwu4 Awaiting requested review from ydwu4 ydwu4 is a code owner

@angelayi angelayi Awaiting requested review from angelayi angelayi is a code owner

Assignees

@XuehaiPan XuehaiPan

Labels
ci-test-showlocals Show local variables on test failures ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request module: dynamo module: pytree open source release notes: dynamo release notes: releng release notes category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

torch.utils._pytree -> stable
4 participants
@XuehaiPan @zou3519 @angelayi @pytorchbot
0