pytorch/torch/distributed/_pycute
fduwjj 232dd65c15 [CuTe] Change the logic of pycute manipulation ops like coalesce, complement 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: https://github.com/pytorch/pytorch/pull/162690
Approved by: https://github.com/ezyang
ghstack dependencies: #162413, #162534, #162414
2025-09-16 19:53:45 +00:00
..
__init__.py [CuTe] Change the logic of pycute manipulation ops like coalesce, complement from co-lex to lex (#162690) 2025-09-16 19:53:45 +00:00
int_tuple.py [CuTe] Change the logic of pycute manipulation ops like coalesce, complement from co-lex to lex (#162690) 2025-09-16 19:53:45 +00:00
layout.py [CuTe] Change the logic of pycute manipulation ops like coalesce, complement from co-lex to lex (#162690) 2025-09-16 19:53:45 +00:00
typing.py