Closed haguettaz closed 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
Thanks also for this one ; I have made the "offset" argument optional, to recover the behavior of numpy.mod by default.
Adds modulo with offset operation on LazyTensors
Test Plan: