Feature Request, adaptive argK

[MH-limarco]

Hi KeOps team, I am working on adaptive KNN, based on probabilistic sparsity mechanism.
Since my sample size N is very large, I really need to improve my original algorithm based on pytorch based on your architecture

import torch
N = 1000 #Sample number 1000~ 100000
dim = 64 # 64 ~ 1500
x = torch.randn(N, dim)
a = torch.nn.Parameter(torch.tensor([0.5])) #parm1 
b = torch.nn.Parameter(torch.tensor([0.5])) #parm2

similarty = torch.cdist(x, x) #Some distance formulas
similarty -= torch.log(-torch.log(torch.rand_like(x) + 1e-6)) #add gumbel_noise

probs = a[0] * torch.relu(similarty - b[0]) # adaptive formulas

res  = (((probs + probs.t()) / 2) > 0) * 1  # Change all values ​​> 0 to 1, else to 0

similarty_idx = probs.nonzero().T # output1, shape(2, L) L is adaptive Neighbor number

loss_probs = probs.sum(dim=1) # output2, shape(N, )

This is My own code with Keops.

from pykeops.torch import LazyTensor
import torch
N = 1000 #Sample number
dim = 64
x = torch.randn(N, dim)
a = torch.nn.Parameter(torch.tensor([0.5])) #parm1
b = torch.nn.Parameter(torch.tensor([0.5])) #parm2

##Some distance formulas
G_i = LazyTensor(x[:, None, :])
X_j = LazyTensor(x[None, :, :])
similarty = ((G_i - X_j) ** 2).sum(-1)

#similarty  -= torch.log(-torch.log(torch.rand_like(similarty) + 1e-6))  # gumbel_noise

# adaptive formulas
probs = a[0] * (similarty - b[0]).relu()
probs = (((probs + probs.t()) / 2) > 0) # Change all values > 0 to 1, else to 0

indices = probs.nonzero().T  # output1  Error: not support nonzero()
loss_probs = probs.sum(dim=1)  # output2

Key points I encountered:
How to add random noise in LazyTensor?
How to get adaptive number K index in keops

I am new to C++/keops, and I wonder if there is any keops method to handle this code?
Thanks!!!~