[dynamo][pytree][1/N] make CXX pytree traceable: `tree_iter` / `tree_leaves` #137397

XuehaiPan · 2024-10-05T11:45:07Z

Stack from ghstack (oldest at bottom):

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

[ghstack-poisoned]

pytorch-bot · 2024-10-05T11:45:11Z

🔗 Helpful Links

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

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

✅ You can merge normally! (4 Unrelated Failures)

As of commit af679d5 with merge base 9012e7a ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

pull / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 4, 5, lf.linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
test_cuda.py::TestCuda::test_randint_generation_for_large_numel
trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 1, 5, lf.linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
test_cuda.py::TestCuda::test_randint_generation_for_large_numel

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#141703)
convnext_base
inductor / cuda12.4-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#141498)
convnext_base

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

…leaves` ghstack-source-id: 5008858 Pull Request resolved: pytorch#137397

[ghstack-poisoned]

…leaves` ghstack-source-id: aa736c9 Pull Request resolved: pytorch#137397

jansel

Could we use the python version of pytree that already exists rather than having a new implementation in polyfils?

[ghstack-poisoned]

…leaves` ghstack-source-id: aa736c9 Pull Request resolved: pytorch#137397

XuehaiPan · 2024-10-06T06:08:37Z

Could we use the python version of pytree that already exists rather than having a new implementation in polyfils?

@jansel We can do this for functions that do not return the treespec variable. E.g.:

generator = tree_iter(tree)
leaves    = tree_leaves(tree)
newtree   = tree_map(func, tree)
tree_map_(func, tree)

I'd implement new polyfills to ensure behavior consistency between polyfill and C++ for the following functions that return/access the treespec variable. See:

[dynamo][pytree][2/N] make CXX pytree traceable: tree_flatten / tree_unflatten / tree_structure #137398

leaves, treespec = tree_flatten(tree)
treespec         = tree_structure(tree)
tree             = tree_unflatten(leaves, treespec)

In most cases, the new polyfill is identical to the already existing Python pytree. For historical reasons, the Python pytree has existed for years and I think splitting the implement makes polyfill maintenance easier. cc @zou3519 for thoughts.

atalman · 2024-12-02T16:02:53Z

@pytorchmergebot revert -c ghfirst -m "Failing internal test"

pytorchmergebot · 2024-12-02T16:05:03Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2024-12-02T16:05:16Z

@XuehaiPan your PR has been successfully reverted.

…/ `tree_leaves` (#137397)" This reverts commit 07850bb. Reverted #137397 on behalf of https://github.com/atalman due to Failing internal test ([comment](#137397 (comment)))

atalman · 2024-12-02T16:57:36Z

HI @XuehaiPan looks like its failing on: https://github.com/pytorch/pytorch/blob/main/test/dynamo/test_trace_rules.py#L328

…leaves` ghstack-source-id: a4960fe Pull Request resolved: pytorch#137397

[ghstack-poisoned]

XuehaiPan · 2024-12-02T18:03:32Z

HI @XuehaiPan looks like its failing on: main/test/dynamo/test_trace_rules.py#L328

Updated.

XuehaiPan · 2024-12-02T18:05:16Z

@jansel @zou3519 Could you take a look at this and the follow-up PRs?

…leaves` ghstack-source-id: 4c6f47b Pull Request resolved: pytorch#137397

[ghstack-poisoned]

…leaves` ghstack-source-id: 8c07273 Pull Request resolved: pytorch#137397

XuehaiPan · 2024-12-03T11:09:57Z

@pytorchbot merge

pytorchmergebot · 2024-12-03T11:11:51Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…leaves` (pytorch#137397) Pull Request resolved: pytorch#137397 Approved by: https://github.com/jansel ghstack dependencies: pytorch#141360

…/ `tree_leaves` (pytorch#137397)" This reverts commit 07850bb. Reverted pytorch#137397 on behalf of https://github.com/atalman due to Failing internal test ([comment](pytorch#137397 (comment)))

…leaves` (pytorch#137397) Pull Request resolved: pytorch#137397 Approved by: https://github.com/jansel

weifengpy · 2024-12-06T23:43:38Z

I am seeing error in latest trunk, likely due to this PR. Is it possible to take a 2nd look if this breaks the trunk?

error about is_dict_insertion_ordered:

ERROR test/distributed/_tensor/test_math_ops.py - ValueError: Duplicate dispatch rule for <built-in method is_dict_insertion_ordered of PyCapsule object at 0x7f343d2b2910>: already registered in VariableBuilder's id dispatch map

how to repro:

from torch.testing._internal.common_utils import run_tests

full error

Traceback (most recent call last):
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/runner.py", line 341, in from_call
    result: Optional[TResult] = func()
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/runner.py", line 372, in <lambda>
    call = CallInfo.from_call(lambda: list(collector.collect()), "collect")
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/python.py", line 531, in collect
    self._inject_setup_module_fixture()
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/python.py", line 545, in _inject_setup_module_fixture
    self.obj, ("setUpModule", "setup_module")
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/python.py", line 310, in obj
    self._obj = obj = self._getobj()
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/python.py", line 528, in _getobj
    return self._importtestmodule()
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/python.py", line 617, in _importtestmodule
    mod = import_path(self.path, mode=importmode, root=self.config.rootpath)
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/site-packages/_pytest/pathlib.py", line 565, in import_path
    importlib.import_module(module_name)
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/data/users/oss/pytorch/test/distributed/_tensor/test_math_ops.py", line 25, in <module>
    from torch.testing._internal.common_utils import run_tests
  File "/data/users/oss/pytorch/torch/testing/_internal/common_utils.py", line 1814, in <module>
    condition=torch._dynamo.config.inline_inbuilt_nn_modules,
  File "/data/users/oss/pytorch/torch/__init__.py", line 2647, in __getattr__
    return importlib.import_module(f".{name}", __name__)
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/data/users/oss/pytorch/torch/_dynamo/__init__.py", line 42, in <module>
    from .polyfills import loader as _  # usort: skip # noqa: F401
  File "/data/users/oss/pytorch/torch/_dynamo/polyfills/loader.py", line 24, in <module>
    POLYFILLED_MODULES: Tuple["ModuleType", ...] = tuple(
  File "/data/users/oss/pytorch/torch/_dynamo/polyfills/loader.py", line 25, in <genexpr>
    importlib.import_module(f".{submodule}", package=polyfills.__name__)
  File "/home/local/miniconda3/envs/torchtitan/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/data/users/oss/pytorch/torch/_dynamo/polyfills/pytree.py", line 31, in <module>
    def _(*args: Any, **kwargs: Any) -> bool:
  File "/data/users/oss/pytorch/torch/_dynamo/decorators.py", line 356, in wrapper
    raise ValueError(
ValueError: Duplicate dispatch rule for <built-in method is_dict_insertion_ordered of PyCapsule object at 0x7f343d2b2910>: already registered in VariableBuilder's id dispatch map
=========================================================================================================== short test summary info ===========================================================================================================
ERROR test/distributed/_tensor/test_math_ops.py - ValueError: Duplicate dispatch rule for <built-in method is_dict_insertion_ordered of PyCapsule object at 0x7f343d2b2910>: already registered in VariableBuilder's id dispatch map

Update

91290da

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels Oct 5, 2024

pytorchbot added the open source label Oct 5, 2024

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

223764b

…leaves` ghstack-source-id: 5008858 Pull Request resolved: pytorch#137397

XuehaiPan added ciflow/trunk Trigger trunk jobs on your pull request module: pytree release notes: dynamo labels Oct 5, 2024

XuehaiPan requested review from zou3519 and jansel October 5, 2024 11:53

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

5eab7ad

…leaves` ghstack-source-id: 5008858 Pull Request resolved: pytorch#137397

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

34c119e

…leaves` ghstack-source-id: 5008858 Pull Request resolved: pytorch#137397

This was referenced Oct 5, 2024

[dynamo][pytree][3/N] make CXX pytree traceable: tree_map / tree_map_ #137399

Closed

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

Open

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

588c7fd

…leaves` ghstack-source-id: 5008858 Pull Request resolved: pytorch#137397

Update

eec7f0c

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

95ed561

…leaves` ghstack-source-id: aa736c9 Pull Request resolved: pytorch#137397

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

96b1334

…leaves` ghstack-source-id: aa736c9 Pull Request resolved: pytorch#137397

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

73dd8ff

…leaves` ghstack-source-id: aa736c9 Pull Request resolved: pytorch#137397

jansel approved these changes Oct 5, 2024

View reviewed changes

jansel requested changes Oct 5, 2024

View reviewed changes

Update

2364bdb

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Oct 6, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

297e098

…leaves` ghstack-source-id: aa736c9 Pull Request resolved: pytorch#137397

XuehaiPan mentioned this pull request Oct 6, 2024

Add XuehaiPan to CODEOWNERS for C++ PyTree utilities #137408

Closed

pytorchmergebot closed this in 07850bb Nov 27, 2024

pytorchmergebot removed the merging label Nov 27, 2024

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Dec 2, 2024

pytorchmergebot reopened this Dec 2, 2024

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Dec 2, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

7a85529

…leaves` ghstack-source-id: a4960fe Pull Request resolved: pytorch#137397

Update

5f6f298

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Dec 2, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

6c21328

…leaves` ghstack-source-id: 4c6f47b Pull Request resolved: pytorch#137397

Update

af679d5

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request Dec 2, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

7a8bcec

…leaves` ghstack-source-id: 8c07273 Pull Request resolved: pytorch#137397

pytorchmergebot added the merging label Dec 3, 2024

pytorchmergebot closed this in 78543e6 Dec 3, 2024

pytorchmergebot removed the merging label Dec 3, 2024

pobin6 pushed a commit to pobin6/pytorch that referenced this pull request Dec 5, 2024

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / `tree_…

20cdaca

…leaves` (pytorch#137397) Pull Request resolved: pytorch#137397 Approved by: https://github.com/jansel

weifengpy mentioned this pull request Dec 6, 2024

AttributeError: 'builtin_function_or_method' object has no attribute '__python_implementation__' #142256

Closed

XuehaiPan mentioned this pull request Dec 18, 2024

[Reland 2.6][dynamo][pytree] make CXX pytree traceable: tree_{flatten,unflatten,structure,map,map_} #143499

Closed

github-actions bot deleted the gh/XuehaiPan/176/head branch January 6, 2025 02:07

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[dynamo][pytree][1/N] make CXX pytree traceable: `tree_iter` / `tree_leaves` #137397

[dynamo][pytree][1/N] make CXX pytree traceable: `tree_iter` / `tree_leaves` #137397

Uh oh!

XuehaiPan commented Oct 5, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Oct 5, 2024 •

edited

Loading

Uh oh!

jansel left a comment

Uh oh!

XuehaiPan commented Oct 6, 2024 •

edited

Loading

Uh oh!

atalman commented Dec 2, 2024

Uh oh!

pytorchmergebot commented Dec 2, 2024

Uh oh!

pytorchmergebot commented Dec 2, 2024

Uh oh!

atalman commented Dec 2, 2024

Uh oh!

XuehaiPan commented Dec 2, 2024

Uh oh!

XuehaiPan commented Dec 2, 2024

Uh oh!

XuehaiPan commented Dec 3, 2024

Uh oh!

pytorchmergebot commented Dec 3, 2024

Uh oh!

weifengpy commented Dec 6, 2024 •

edited

Loading

Uh oh!

Uh oh!

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / tree_leaves #137397

[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / tree_leaves #137397

Uh oh!

Conversation

XuehaiPan commented Oct 5, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 5, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (4 Unrelated Failures)

Uh oh!

jansel left a comment

Choose a reason for hiding this comment

Uh oh!

XuehaiPan commented Oct 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

atalman commented Dec 2, 2024

Uh oh!

pytorchmergebot commented Dec 2, 2024

Uh oh!

pytorchmergebot commented Dec 2, 2024

Uh oh!

atalman commented Dec 2, 2024

Uh oh!

XuehaiPan commented Dec 2, 2024

Uh oh!

XuehaiPan commented Dec 2, 2024

Uh oh!

XuehaiPan commented Dec 3, 2024

Uh oh!

pytorchmergebot commented Dec 3, 2024

Merge started

Uh oh!

weifengpy commented Dec 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

[dynamo][pytree][1/N] make CXX pytree traceable: `tree_iter` / `tree_leaves` #137397

[dynamo][pytree][1/N] make CXX pytree traceable: `tree_iter` / `tree_leaves` #137397

XuehaiPan commented Oct 5, 2024 •

edited

Loading

pytorch-bot bot commented Oct 5, 2024 •

edited

Loading

XuehaiPan commented Oct 6, 2024 •

edited

Loading

weifengpy commented Dec 6, 2024 •

edited

Loading