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 Batched Select for devices and tensor_ops #182

Merged
merged 14 commits into from
Sep 18, 2022
Merged

Conversation

coreylowman
Copy link
Owner

Closes #177

Unblocks #121

Batched select needs the source type to be broadcasted one level. It was actually impossible to specify this with the macro and axis based select previously used in device. This is because broadcasted select doesn't select from an axis at first, so you can't specify a single thing for it to select from.

Summary:

  • Reworks the devices implementation to use a recursive trait implementation using recursive types to deconflict implementation.
    • These are controlled via unit structs called "Modes"
  • Adds (private) generic select function that uses select modes to control behavior
  • Rework existing tensor impls to use modes
  • Adds a SelectBatchAx0 trait with select_batch function for 1d/2d/3d tensors

src/devices/select.rs Outdated Show resolved Hide resolved
src/devices/select.rs Outdated Show resolved Hide resolved
src/devices/select.rs Outdated Show resolved Hide resolved
src/devices/select.rs Outdated Show resolved Hide resolved
@coreylowman coreylowman merged commit 133e2d6 into main Sep 18, 2022
@coreylowman coreylowman deleted the batched-select branch September 18, 2022 18:34
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.

Add batched select
1 participant