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

Added TryUnstack for tensors. #919

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

Conversation

swfsql
Copy link
Contributor

@swfsql swfsql commented Feb 6, 2024

  • Adds the try_/unstack method for tensors.
  • Also added from_fn for Arrays.
  • This is similar to Adds unstack #831 by @emchristiansen - thanks! I used your code for reference - but the difference is that this one also allows for unstacking into arrays instead of only unstacking into Vecs, which is closer to how the interface for stack currently works.
    • But note that this is only useful for "shape certainty" purposes, there shouldn't be much difference in performance.
    • Also, in this implementation, a two-pass array/vec construction is necessary, which does have a performance hit even for the Vec case.
      • This should be avoidable by having a method allowing to merge a tape into a tensor that already had a OwnedTape (instead of moving the tensor - which is necessary since the type is going to change - and splitting it), but in order to avoid adding such functionality I opted for a two-pass array/vec construction.
  • Together with Add split_tensor_along method #917, may close Please add a split method, to iterate over subtensors of the given tensor #825.

Remaining Tasks

  • Decide on how to manage non-contiguous data - it's mandatory. Possibly panic at runtime.
  • Add Cuda kernel.
  • Consider adding function to merge a OwnedTape into a tensor which already has a OwnedTape, with a &mut to the tensor.
  • Consider adding a layer, even though TryStack doesn't have one yet.
  • Consider adding a _vec method that always realizes the first dimension into a usize before unstacking.
  • Consider allowing the unstack for an arbitrary axis, similar to the concat interface.
    • Note: This is not present on the stack interface; Maybe this can be possible with a helper function that always transmutes that axis as the first one and then calls the unstack; Also this could be made by a future pr.

@swfsql swfsql mentioned this pull request Feb 7, 2024
13 tasks
- Also added `from_fn` for Arrays.

Note: the interface currently requires two passes for construction, one for creating a list of tensors with NoneTape and another for putting tapes into those tensors.
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.

Please add a split method, to iterate over subtensors of the given tensor
2 participants