getkeops / keops

KErnel OPerationS, on CPUs and GPUs, with autodiff and without memory overflows
https://www.kernel-operations.io
MIT License
1.04k stars 63 forks source link

Modulo with offset operation on LazyTensors #136

Closed haguettaz closed 3 years ago

haguettaz commented 3 years ago

Adds modulo with offset operation on LazyTensors

Test Plan:

import pykeops
import torch
import math
from pykeops.torch import LazyTensor

def torch_mod(input, modulus, offset):
    return input - modulus * torch.floor((input - offset)/modulus)

device = 'cuda' 

x = torch.rand(1000, 1)*2*math.pi
y = x.data.clone()
x = x.to(device)
y = y.to(device)
x.requires_grad = True
y.requires_grad = True

x_i = LazyTensor(x[:, None])
s1 = x_i.mod(math.pi, -math.pi/2).sum(0)
s2 = torch.sum(torch_mod(y, math.pi, -math.pi/2))
print("s1 - s2", torch.abs(s1 - s2).item())
assert torch.abs(s1 - s2) < 1e-3, torch.abs(s1 - s2)

s1.backward()
s2.backward()

print("grad_s1 - grad_s2", torch.max(torch.abs(x.grad - y.grad)).item())
assert torch.max(torch.abs(x.grad - y.grad)) < 1e-3
joanglaunes commented 3 years ago

Thanks also for this one ; I have made the "offset" argument optional, to recover the behavior of numpy.mod by default.