Error: incompatible dimensions in Kmin_argKmin

[psclklnk]

Hello everyone,

First of all, thank you for this amazing library! It allowed me to speed up my computations in so many applications already, and I am sure it will do so in the future, too!

While using the library, I came across an error in the Kmin_argKmin operation that seems to be triggered by enabling chunking (similar to this issue). Here is a minimal example that provokes the error on my machine using pykeops 2.1.2:

from pykeops.torch import LazyTensor

n1 = 1
n2 = 3096
dim= 300
lt_x = LazyTensor(torch.rand((n1, 1, dim), device="cuda:0"))
lt_y = LazyTensor(torch.rand((1, n2, dim), device="cuda:0"))
squared_dist = ((lt_x - lt_y) ** 2).sum(-1)
res = (-squared_dist).Kmin_argKmin(2, dim=1, enable_chunks=True)

Running the same script with enable_chunks=False works just fine. The only downside to this solution is that it runs about 5 times slower for my problem. Currently, I simply create a workaround by catching any errors in Kmin_argKmin and re-running the computation with chunking disabled.

I did not check the codebase for a fix and hope that the problem can be fixed similarly to the one I linked above.

Best Pascal

[joanglaunes]

Hello @psclklnk , I just tried now and it looks that the fix from issue #294 fixes it in your case also. So your code works with the current main branch of pykeops, and the only problem is that we have not done any release since quite a while... We will try to do a release as soon as possible, sorry!