Comparison operators for LazyTensor

[Louis-Pujol]

Hi !

I wonder if it would be possible to add support for comparison operators between LazyTensors and numbers. The following code :

from pykeops.torch import LazyTensor
import torch

x = torch.rand(10, 3)
y = torch.rand(10, 3)

x_i = LazyTensor(x[:, None, :])
y_j = LazyTensor(y[None, :, :])

threshold = 0.1 ** 2

D_ij = ((x_i - y_j) ** 2)

D_ij = D_ij <= threshold

leads to the error message TypeError: '<=' not supported between instances of 'LazyTensor' and 'float', while D_ij = (threshold - D_ij).step() do what was expected (D_ij acts as a boolean tensor, encoded with 0s and 1s).

Thanks,

[joanglaunes]

Hello @Louis-Pujol , This is now done in the main branch. I have added <, >, <=, >=, == and != ; it works between two instances of LazyTensor or one LazyTensor and a int or float.

[Louis-Pujol]

Thanks @joanglaunes !