Add split_tensor_along method #917
Draft
Conversation
This file contains hidden or 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
This was referenced Feb 2, 2024
Closed
Remove ftz
Avoid ci errors
- Add `TrySplitShapeAlong` and `TrySplitTensorAlong`. - Minor linting and docs fix. TODO - Check if the tape should be returned. If not, it can be removed from the interface. - Add cuda kernel. - Consider a different interface, where it could get split in more than two tensors - possibly stated on a vec. In this way it could get closer to the pytorch interface (chunks).
693b699 to
aa44eb5
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
TrySplitShapeAlongandTrySplitTensorAlong.This is a module that allows a tensor to be split in two - can think of it as the opposite of what the concat does - and the concat was the reference for the implementation.
I'm unsure on how to deal with the tape. I had the impression that it would be needed to be returned and merged with one of the chunks - the idea of returning it would mean that it wouldn't be fixed to eg. get merged to the first chunk - but the backpropagation tests appears to not require this merging of tapes?
Remaining tasks:
Note: I haven't made an actual pytorch test to compare, not even locally, but have made a bigger test that depend on this functionality and it appeared to work ok.