Genred sparse reduction fails silently for int64 indices

[Turakar]

I try to use sparsity with Genred, but the code fails silently.

import pykeops.torch.cluster
import torch
from pykeops.torch import Genred

def main():
    x = torch.linspace(0, 1, 4).unsqueeze(-1)
    y = torch.linspace(0, 1, 4).unsqueeze(-1)

    mvm = Genred(
        "x * y",
        ["x = Vi(1)", "y = Vj(1)"],
        reduction_op="Sum",
        axis=1,
    )

    ranges1 = torch.tensor([[0, 2], [2, 4]])
    ranges2 = torch.tensor([[0, 2], [2, 4]])
    keep = torch.tensor([[True, True], [True, True]])
    ranges = pykeops.torch.cluster.from_matrix(ranges1, ranges2, keep)

    print("\n".join([str(x) for x in ranges]))

    result = mvm(x, y, ranges=ranges)

    reference = torch.sum(x[:, None] * y[None, :], dim=-2)
    print(result)
    print(reference)

if __name__ == "__main__":
    main()

Produces the following output:

/path/to/venv/lib64/python3.11/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3483.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
tensor([[0, 2],
        [2, 4]])
tensor([2, 4], dtype=torch.int32)
tensor([[0, 2],
        [2, 4],
        [0, 2],
        [2, 4]])
tensor([[0, 2],
        [2, 4]])
tensor([2, 4], dtype=torch.int32)
tensor([[0, 2],
        [2, 4],
        [0, 2],
        [2, 4]])
tensor([[0.],
        [0.],
        [0.],
        [0.]])
tensor([[-3.8142],
        [ 0.9167],
        [-1.1071],
        [-5.8596]])

As you can see, the result is all zeros, which is unequal from the reference calculation. I would expect the last two outputs to be equal. Is there a mistake in my example or is there a bug in KeOps?

Versions:

• Python 3.11.3
• PyTorch 2.0.0+cu117
• CUDA 11

The reference example seems to work in my environment.

[Turakar]

Seems like this error can be prevented by casting the indices to int32 before passing them to from_matrix().