Wrong output with batch dimensions on parameters

[joanglaunes]

When using batch dimensions on parameter variables, the computation is performed only on the first row. Here is a simple example with gaussian kernel convolution:

import torch
from pykeops.torch import LazyTensor, Genred

device = "cuda" if torch.cuda.is_available() else "cpu"

B, M, N = 3, 5, 4

x = LazyTensor(torch.rand(1, M, 1, 1, device=device))
y = LazyTensor(torch.rand(1, 1, N, 1, device=device))
p = LazyTensor(torch.rand(B, 1, 1, 1, device=device))

res = (-p*(x-y)**2).exp().sum(dim=2)

print(res.reshape((B,M)))

Output gives

tensor([[3.9838, 3.9832, 3.9790, 3.9049, 3.8810],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]])

This happens both in cpu and gpu modes.
This is due I think to the way we handle size of the output in the code. There is a confusion with the variable nx that sometimes refers to the size of the output (B*M in the previous example) and sometimes to the size of the i variables (M).

[bcharlier]

It seems to be related to broadcasting:

import torch
from pykeops.torch import LazyTensor, Genred

device = "cuda" if torch.cuda.is_available() else "cpu"

# No broadcasting : WORKS
B, M, N = 3, 5, 4
_x = torch.rand(B, M, 1, 1, device=device)
x = LazyTensor(_x)
_p = torch.rand(B, 1, 1, 1, device=device)
p = LazyTensor(_p)


print(
    torch.allclose(
        (p * x).sum(dim=2).reshape(B, M),
        (_p * _x).sum(dim=2).reshape(B, M)
    )
) 
# True

while

import torch
from pykeops.torch import LazyTensor, Genred

device = "cuda" if torch.cuda.is_available() else "cpu"


# Broadcasting : DOES NOT WORK
B, M, N = 3, 5, 4
_x = torch.rand(1, M, 1, 1, device=device)
x = LazyTensor(_x)
_p = torch.rand(B, 1, 1, 1, device=device)
p = LazyTensor(_p)


print((p * x).sum(dim=2).reshape(B, M))
print((_p * _x).sum(dim=2).reshape(B, M))

gives

tensor([[0.0400, 0.4201, 0.1151, 0.2437, 0.1425],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]], device='cuda:0')
tensor([[0.0400, 0.4201, 0.1151, 0.2437, 0.1425],
        [0.0113, 0.1189, 0.0326, 0.0690, 0.0403],
        [0.0281, 0.2947, 0.0807, 0.1710, 0.1000]], device='cuda:0')