Add some shape decomps (t, transpose, rot90, stack) #76873
Conversation
🔗 Helpful links
💊 CI failures summary and remediationsAs of commit 8a03611 (more details on the Dr. CI page): Expand to see more💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Please report bugs/suggestions to the (internal) Dr. CI Users group. |
Chillee
changed the title
torch/_decomp/decompositions.py
Outdated
| assert dims[0] != dims[1] and abs(dims[0] - dims[1]) != total_dims, f"expected rotation dims to be different, but got dim0 = {dims[0]} and dim1 = {dims[1]}" | ||
| assert dims[0] < total_dims and dims[0] >= -total_dims, f"Rotation dim0 out of range, dim0 = {dims[0]}" | ||
| assert dims[1] < total_dims and dims[1] >= -total_dims, f"Rotation dim1 out of range, dim1 = {dims[1]}" | ||
| k = (4 + (k % 4)) % 4 |
There was a problem hiding this comment.
what's this? Modulo of a negative number is what you'd expect (so k=-1 will translate to 3)
There was a problem hiding this comment.
uh... too literal translation of C++ code to Python :P
|
|
||
| if idx < 0: | ||
| _idx = idx + _rank | ||
| else: |
|
@pytorchbot merge this please |
|
(passed all tests on a previous commit and I just fixed a merge issue, so hopefully good) :) |
|
Hey @Chillee. |
Chillee
added
module: internals
Summary: Also fixes xlogy (turns out the only thing it was missing was a type cast annotation! nice!) I also renamed `canonicalize_idx` => `canonicalize_dim` (to align with `canonicalize_dims`) and fixed a bug in it (cc: mruberry) Pull Request resolved: #76873 Approved by: https://github.com/mruberry Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/e9f34931ef5a4172eac1cc2238218b94beb59b28 Reviewed By: malfet Differential Revision: D36201593 fbshipit-source-id: 6defbeb0b95ecf205211ed1e2817171905d4c84c
Also fixes xlogy (turns out the only thing it was missing was a type cast annotation! nice!)
I also renamed
canonicalize_idx=>canonicalize_dim(to align withcanonicalize_dims) and fixed a bug in it (cc: @mruberry)