Tensor Filtering

A common operation for me is filtering down a tensor to a subset of it based on some condition. For example, I may want to filter out all the rows of a tensor whose values are all NaN. What would be helpful here is, e.g., a `.indexes_where_one()` function that takes a 1-D tensor of boolean (0/1) values and returns a tensor on the same device of all the indexes that are one. This can then be fed into other functions like `.index_select()` to do filtering. Currently, I do this with my own version of the function that copies to the CPU, but it would be great to have a function built into Candle and that doesn't have to copy between devices.