Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add split_tensor_along method #917

Draft
wants to merge 8 commits into
base: main
Choose a base branch
from

Conversation

swfsql
Copy link
Contributor

@swfsql swfsql commented Feb 1, 2024

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:

  • 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).
  • Add a layer in the similar vein to the concat one.

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.

This was referenced Feb 2, 2024
- 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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants