getkeops / keops

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

Arg at position X is not contiguous in generic reduction #245

Closed fwilliams closed 2 years ago

fwilliams commented 2 years ago

Hi,

I'm using pykeops.torch.Genred with a formula string and list of aliases to evaluate a matrix/vector product. My matrix expression is somewhat complex and the full matvec needs to be evaluated in parts.

Unfortunately, the backward pass of the reduction is giving the error:

RuntimeError: [Keops] Arg at position 10: is not contiguous. Please provide 'contiguous' dara array, as KeOps does not support strides. If you're getting this error in the 'backward' pass of a code using torch.sum() on the output of a KeOps routine, you should consider replacing 'a.sum()' with 'torch.dot(a.view(-1), torch.ones_like(a).view(-1))'. 

The backward reduction is:

[pyKeOps] Compiling libKeOpstorch0b9260d65a in /home/fwilliams/.cache/pykeops-1.5-cpython-39:
       formula: Grad_WithSavedForward(Sum_Reduction(Exp(-half * SqDist(xp, yp)) * (dyf_dxpi | xf) + (xf | yf) * (-(xpi - ypi) * Exp(-half * SqDist(xp, yp)) / h) * b,0), Var(1,64,0), Var(10,1,0), Var(11,1,0))
       aliases: xp = Vi(0,3); xf = Vi(1,64); xpi = Vi(2,1); yp = Vj(3,3); yf = Vj(4,64); ypi = Vj(5,1); dyf_dxpi = Vj(6,64); b = Vj(7,1); half = Pm(8,1); h = Pm(9,1); Var(10,1,0); Var(11,1,0); 
       dtype  : float64

The suggestion in the error message (to use a dot product with ones()) does not seem directly applicable here (there isn't a Dot reduction for Genred AFAIK). Is there a way to fix this without rewriting my code using lazy tensors?

fwilliams commented 2 years ago

I misunderstood the error message, sorry. The issue is that I was stacking the output for each block and computing a loss on the stacked tensor.