[Question] How to perform reduction on 4D tensor

[mh-nguyen712]

Hi and Thanks for creating this library!
I have a problem and pykeops seems to be a good frame work.
I have 2 tensors x of shape B, M, D and y of shape B, N, D and I want to compute the distance matrix between these two tensors, resulting in a tensor of shape B, M, N.
The usual code in pytorch is

import torch 
B = 128
M = 1024
N = 1024
D = 1024

x = torch.randn(B, M, D, device='cuda')
y = torch.randn(B, N, D, device='cuda')

dist = (x.unsqueeze(2) - y.unsqueeze(1)).pow(2).sum(dim = -1) 

I've tried using pykeops:

a_i = LazyTensor(x.unsqueeze(1))          # shape (7, 1, 5, 3)
b_j = LazyTensor(y.unsqueeze(2))          # shape (7, 5, 1, 3)
c = (a_i - b_j).sum(dim=-1)

but c is still a LazyTensor. How could I perform the computation?
Thanks for your help!

[joanglaunes]

Hello @mh-nguyen712 ,
Indeed it is a normal behaviour : you are performing a sum over the last axis, which for KeOps means you are still doing an internal virtual operation and not a reduction. Hence the resulting tensor is still a LazyTensor. You would trigger the computation and get an actual tensor by summing over axis=1 or axis=2.
A distance matrix is typically what we want to keep as a virtual tensor in KeOps because otherwise it would have a quadratic memory footprint. Now everything depends on what you want to do with this distance matrix. If it is the final output for your problem, KeOps is useless. But if you want to further compute something from it, KeOps can probably be used.