Fix backward for reshape() on jagged layout NT #117137
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/117137
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 34502fc with merge base 81b7a09 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
| if func is torch._C._nn.scaled_dot_product_attention: | ||
| return jagged_scaled_dot_product_attention(*args, **kwargs) | ||
|
|
||
| # Handle reshape() / reshape_as() here because they're CompositeImplicit. |
There was a problem hiding this comment.
NB: Handling these in __torch_function__() this isn't enough to get them working for autograd formulas.
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see Details for Dev Infra teamRaised by workflow job |
Provides symbolic C++-side `reshape_as()` / `reshape()` decomps for jagged layout NTs to make the backwards pass work. [ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour 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 |
| return std::make_pair(M, N); | ||
| } | ||
|
|
||
| Tensor reshape_nested(const Tensor& self, IntArrayRef proposed_shape); |
There was a problem hiding this comment.
Noob question. Why we do need to add this now? And it seems that reshape_nested is not used anywhere after you changed it to reshape_nested_symint in native_functions.yaml?
There was a problem hiding this comment.
Sure thing - good question! This was originally added for strided (C++) NTs, and we don't support nt.sizes() on those. To expand this to support jagged layout NTs, we need to operate in symbolic shapes land, as we always expect them to have a symbolic shape (i.e. a SingletonSymInt for the jagged dimension).
And it seems that reshape_nested is not used anywhere after you changed it to reshape_nested_symint in native_functions.yaml?
Note that reshape_nested_symint() calls the old reshape_nested() for the case of strided NT.
There was a problem hiding this comment.
I see, thanks. Totally missed the usage inside reshape_nested_symint.
I am wondering why do we need to define "reshape_nested" in the header file while most other functions don't? Does it make any difference?
There was a problem hiding this comment.
I am wondering why do we need to define "reshape_nested" in the header file while most other functions don't? Does it make any difference?
The reason for this is because reshape_nested() is no longer referenced from native_functions.yaml. Functions in the dispatch table there will automatically have a header signature generated.
4 participants
Stack from ghstack (oldest at bottom):
Provides symbolic C++-side
reshape_as()/reshape()decomps for jagged layout NTs to make the backwards pass work.