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

boolean indexing ops: AtenNonzeroOp, AtenIndexTensorOp, AtenMaskedSelectOp #3293

Open
Tracked by #347
renxida opened this issue May 6, 2024 · 0 comments
Open
Tracked by #347
Assignees

Comments

@renxida
Copy link
Collaborator

renxida commented May 6, 2024

Ops that need work:

AtenNonzeroOp, // creates a list of indices where a given tensor is nonzero
AtenIndexTensorOp, // currently only supports constant shape inputs / outputs & may not support boolean indices
AtenMaskedSelectOp, // no lowering. similar behavior to indexing with a boolean, but [flattens the output](https://discuss.pytorch.org/t/masked-select-vs-indexing-with-boolean-in-aten/33621/3)

All of the above ops decompose to indexing with a Boolean tensor

The key thing I can't figure out for these ops right now is that I really need a way to shorten something like

[0, 0, 1, 2, 0, 5, 8]

to

[1,2,5,8]

some kind of list filtering, basically.

If I have any one of the above ops, I would be able to decompose the rest to it. Indexing seems the most canonical choice.

How to implement bool indexing

I'm thinking just using a scf for loop to append to a list, then convert to a variable-length tensor. There are some nuances there and I should probably read this discussion with Ben and Quinn to figure it out https://discord.com/channels/973663919757492264/1172976577924050984/1220489866161295440

@renxida renxida self-assigned this May 6, 2024
@renxida renxida changed the title [aten to linalg] Boolean indexing ops [TORCH] Boolean indexing ops May 6, 2024
@renxida renxida changed the title [TORCH] Boolean indexing ops [torch] Boolean indexing ops May 6, 2024
@renxida renxida changed the title [torch] Boolean indexing ops boolean indexing ops: AtenNonzeroOp, AtenIndexTensorOp, AtenMaskedSelectOp May 6, 2024
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

No branches or pull requests

1 participant