[CuTe] Add type for CuTe layout via claude by fduwjj · Pull Request #162534 · pytorch/pytorch

fduwjj · 2025-09-09T22:22:44Z

Stack from ghstack (oldest at bottom):

This PR mostly is a cosmetic change using Claude to add types for copied PyCute code.
We removed all suppressions of linters and add type checker, type alias and mypy ignore(if needed) so that the pycute code will be checked by linter.

cc @H-Huang @awgu @wanchaol @fegin @wz337 @wconstab @d4l3k @pragupta @ezyang @msaroufim @dcci

[ghstack-poisoned]

pytorch-bot · 2025-09-09T22:22:48Z

🔗 Helpful Links

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

📄 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

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

✅ No Failures

As of commit 0428777 with merge base 5babb4d ():
💚 Looks good so far! There are no failures yet. 💚

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

ghstack-source-id: 2a11926 Pull Request resolved: #162534

This PR mostly is a cosmetic change using Claude to add types for copied PyCute code. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

ghstack-source-id: 6cba43a Pull Request resolved: #162534

torch/distributed/_pycute/int_tuple.py

torch/distributed/_pycute/layout.py

ezyang · 2025-09-10T03:09:35Z

You should be able to feed my CR comments directly to Claude. Please review the output yourself though.

This PR mostly is a cosmetic change using Claude to add types for copied PyCute code. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

ghstack-source-id: 0cec363 Pull Request resolved: #162534

torch/distributed/_pycute/int_tuple.py

torch/distributed/_pycute/layout.py

torch/distributed/_pycute/int_tuple.py

This PR mostly is a cosmetic change using Claude to add types for copied PyCute code. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

ghstack-source-id: 9ed51e9 Pull Request resolved: #162534

torch/distributed/_pycute/int_tuple.py

This PR mostly is a cosmetic change using Claude to add types for copied PyCute code. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]

fduwjj · 2025-09-12T04:51:52Z

@pytorchbot merge

pytorchmergebot · 2025-09-12T04:53:35Z

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

… DeviceMesh (#162414) We create a wrapper class acting as a layout for device mesh so that we can add new methods more specific to DeviceMesh and keep the core logic of CuTe manipulation inside pycute module. This PR create the main body of the code and then next PR will come with actual implementation and unit test for device mesh layout. (Actual implementation can be found in #161016) Pull Request resolved: #162414 Approved by: https://github.com/ezyang ghstack dependencies: #162413, #162534

… DeviceMesh (#162414) We create a wrapper class named "_MeshLayout" acting as a layout for device mesh so that we can add new methods more specific to DeviceMesh and keep the core logic of CuTe manipulation inside pycute module. This PR create the main body of the code and then next PR will come with actual implementation and unit test for device mesh layout. (Actual implementation can be found in #161016) Pull Request resolved: #162414 Approved by: https://github.com/ezyang, https://github.com/fegin ghstack dependencies: #162413, #162534

…plement from co-lex to lex (#162690) PyTorch tensor iteration (.view, contiguous, broadcasting) and NumPy array indexing all follow lexicographic (row-major) order. In Lexicographic (lex) on (i0, i1, …, i{k-1}): the leftmost index(stride is larger) changes fastest and the rightmost index changes slowest and usually last dim is contiguous. However original pycute is all based on co-lex, after porting their code into pytorch and some cosmetic change, we now make it lex so that we can use it for use cases like device mesh internal bookkeeping and other stuff as well. Changes included in this PR: 1. We changes all API ported in, included prefix_product(stride inferring and rename it to suffix_product), idx2crd, crd2idx, coalesce, composition, complement, right_inverse and left_inverse to make sure they are working in the lex way. 2. Added more unit test cases for some API mentioned above since existing unit tests do not have full coverage. 3. One bug fix inside composition, which will lead to infinite recursive call. Pull Request resolved: #162690 Approved by: https://github.com/ezyang ghstack dependencies: #162413, #162534, #162414

This PR mostly is a cosmetic change using Claude to add types for copied PyCute code. We removed all suppressions of linters and add type checker, type alias and mypy ignore(if needed) so that the pycute code will be checked by linter. Pull Request resolved: pytorch#162534 Approved by: https://github.com/ezyang, https://github.com/Skylion007 ghstack dependencies: pytorch#162413

… DeviceMesh (pytorch#162414) We create a wrapper class acting as a layout for device mesh so that we can add new methods more specific to DeviceMesh and keep the core logic of CuTe manipulation inside pycute module. This PR create the main body of the code and then next PR will come with actual implementation and unit test for device mesh layout. (Actual implementation can be found in pytorch#161016) Pull Request resolved: pytorch#162414 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#162413, pytorch#162534

… DeviceMesh (pytorch#162414) We create a wrapper class named "_MeshLayout" acting as a layout for device mesh so that we can add new methods more specific to DeviceMesh and keep the core logic of CuTe manipulation inside pycute module. This PR create the main body of the code and then next PR will come with actual implementation and unit test for device mesh layout. (Actual implementation can be found in pytorch#161016) Pull Request resolved: pytorch#162414 Approved by: https://github.com/ezyang, https://github.com/fegin ghstack dependencies: pytorch#162413, pytorch#162534

…plement from co-lex to lex (pytorch#162690) PyTorch tensor iteration (.view, contiguous, broadcasting) and NumPy array indexing all follow lexicographic (row-major) order. In Lexicographic (lex) on (i0, i1, …, i{k-1}): the leftmost index(stride is larger) changes fastest and the rightmost index changes slowest and usually last dim is contiguous. However original pycute is all based on co-lex, after porting their code into pytorch and some cosmetic change, we now make it lex so that we can use it for use cases like device mesh internal bookkeeping and other stuff as well. Changes included in this PR: 1. We changes all API ported in, included prefix_product(stride inferring and rename it to suffix_product), idx2crd, crd2idx, coalesce, composition, complement, right_inverse and left_inverse to make sure they are working in the lex way. 2. Added more unit test cases for some API mentioned above since existing unit tests do not have full coverage. 3. One bug fix inside composition, which will lead to infinite recursive call. Pull Request resolved: pytorch#162690 Approved by: https://github.com/ezyang ghstack dependencies: pytorch#162413, pytorch#162534, pytorch#162414