Make fbgemm::jagged_index_select pt2_compliant #2170
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
✅ Deploy Preview for pytorch-fbgemm-docs canceled.
|
|
This pull request was exported from Phabricator. Differential Revision: D51670069 |
Summary: It needed an abstract impl. Reviewed By: williamwen42 Differential Revision: D51647394
Summary: See title. Reviewed By: williamwen42 Differential Revision: D51647393
Summary: - Added abstract impl in sparse_ops.py. I don't think it's worth splitting up the abstract impls into multiple .py files right now, unless someone comes to us with memory issues. Reviewed By: williamwen42 Differential Revision: D51647392
Summary: Also deleted two skips that were marked as flaky. Those don't appear to actually be flaky. Reviewed By: williamwen42 Differential Revision: D51647391
Summary: - The problem with the original op was that in the Autograd implementation, it needed to call Tensor.item(). This doesn't work with FakeTensors (maybe it can some day in the future). - We create two new ops, `jagged_index_select_2d_forward_v2` and `jagged_index_add_2d_forward_v2` (which is effectively the backward) that do the Tensor.item() calls, and change fbgemm::jagged_index_select's Autograd implementation to call those. - We add abstract impls for those two new ops. - Finally, we move the fbgemm::jagged_index_select implementation to CompositeImplicitAutograd (and delete the CPU/CUDA impls, because those are redundant). Differential Revision: D51670069
zou3519
force-pushed
the
export-D51670069
branch
from
a0ecccb
to
a5e0137
Compare
|
This pull request was exported from Phabricator. Differential Revision: D51670069 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D51670069 |
|
This pull request has been merged in 1c40928. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
needed to call Tensor.item(). This doesn't work with FakeTensors (maybe it
can some day in the future).
jagged_index_select_2d_forward_v2andjagged_index_add_2d_forward_v2(which is effectively the backward) that dothe Tensor.item() calls, and change fbgemm::jagged_index_select's Autograd
implementation to call those.
CompositeImplicitAutograd (and delete the CPU/CUDA impls, because those are
redundant).
Differential Revision: D51670069