boolean indexing ops: AtenNonzeroOp, AtenIndexTensorOp, AtenMaskedSelectOp

[renxida]

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

[vivekkhandelwal1]

Hi @renxida, are you still working on these ops?

[renxida]

Abandoned them before becasue thought they weren't important.

Just got a message from Dhiraj that these are important. Hopping on them.

[renxida]

using some logic from 33f9101

[renxida]

some prototyping work. i implemented the decomposition in python and tested them against pytorch

https://colab.research.google.com/drive/1j724QyrwrZojEoBnng6w87LADLGvlb6i#scrollTo=3vEqg6uopswY

and this code exports the ir for my prototypes:
https://gist.github.com/renxida/87b4e9e68e547ab845e717e91208b4d1

i had to use torchscript because fx doesn't allow slicing with tensor_a[:sum(tensor_b). this would probably be a problem later.