Feature Request, getting indices from a LazyTensor

[yaadhavraajagility]

I have a specific use case where i have a [B, N, M] tensor that stores distance values.
I would like to do an operation akin to sample_from_tensor[squared_distance_lazy_tensor > radius]

But it looks like TypeError: '>' not supported between instances of 'LazyTensor' and 'LazyTensor'

Alternatively, it looks like you support Kmin_argKmin. Instead of the k nearest neighbour approach, do you have something that can return indices say based on distance?

[joanglaunes]

Hello @yaadhavraajagility ,
About your first question, comparison operators have been added recently into KeOps, so it is available when using the main git branch, but not yet in the released version of KeOps.
Now about your second question, it is not clear what you would like to get exactly. If it is for example to get all indices of points whose distance from the current point is below some threshold, then the problem is that this is has quadratic memory footprint in general, similar to the distance matrix itself, and KeOps currently does not return O(MN) arrays since its purpose is to avoid them. Then if the purpose is to sample from these selected points as in your first question, then using the > comparison operation will give a binary mask (as a LazyTensor), and then you can perform a reduction operation which will correspond to the sampling, so there should be no need for indices.
So could you tell us what would be your use case for returning indices ?