Make fbgemm::jagged_index_select pt2_compliant #2170
Closed
Commits on Nov 29, 2023
-
Make fbgemm::masked_select_jagged_1d pt2_compliant (pytorch#2174)
Summary: It needed an abstract impl. Reviewed By: williamwen42 Differential Revision: D51647394
-
Add generate_opcheck_tests to input_combine_test (pytorch#2173)
Summary: See title. Reviewed By: williamwen42 Differential Revision: D51647393
-
Make fbgemm::tbe_input_combine pt2_compliant (pytorch#2172)
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
-
Mark some more ops as pt2_compliant_tag (pytorch#2171)
Summary: Also deleted two skips that were marked as flaky. Those don't appear to actually be flaky. Reviewed By: williamwen42 Differential Revision: D51647391
-
Make fbgemm::jagged_index_select pt2_compliant (pytorch#2170)
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
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.