search

getkeops / keops

KErnel OPerationS, on CPUs and GPUs, with autodiff and without memory overflows
https://www.kernel-operations.io
MIT License
1.02k stars 65 forks source link

Feature Request, getting indices from a LazyTensor #341

Open yaadhavraajagility opened 8 months ago

yaadhavraajagility commented 8 months ago

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 commented 7 months ago

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 ?