segfault in specific formula

[albertfgu]

I upgraded to the most recent packages (I think keopscore got an update?). One of my keops formulas works great, but the other one is segfaulting with no error message. This is the kernel

def log_vandermonde(v, x, L, conj=True):
    expr = 'ComplexMult(v, ComplexExp(ComplexMult(x, l)))'
    vandermonde_mult = Genred(
        expr,
        [
            'v = Vj(2)',
            'x = Vj(2)',
            'l = Vi(2)',
        ],
        reduction_op='Sum',
        axis=1,
    )

    l = torch.arange(L).to(x)
    v, x, l = _broadcast_dims(v, x, l)
    v = _c2r(v)
    x = _c2r(x)
    l = _c2r(l)

    r = vandermonde_mult(v, x, l, backend='GPU')

[joanglaunes]

Hello @albertfgu ,
Sorry for being very late to reply, I hope my answer will still be of interest for you. I tried your function and it appears to work. But it will fail if L is large and you define the input x to be random complex numbers, because multiplication by integers in the range 0 to L will give complex numbers with large absolute values, and then passing to the exponential will give infinite values. So maybe this was the cause in your case, if you just tried the function without being careful about the x input ? If you define x to be pure imaginary numbers, as is done in the code below, then there should be no problem.

Here is the full code that I tried and which works :

import torch
from pykeops.torch import Genred

def _broadcast_dims(*args):
    out = []
    for x in args:
        out.append(x[:,None])
    return out

def _c2r(x):
    return torch.stack((torch.real(x),torch.imag(x)),dim=len(x.shape)).reshape(x.shape[:-1]+(-1,))

def log_vandermonde(v, x, L, conj=True):
    expr = 'ComplexMult(v, ComplexExp(ComplexMult(x, l)))'
    vandermonde_mult = Genred(
        expr,
        [
            'v = Vj(2)',
            'x = Vj(2)',
            'l = Vi(2)',
        ],
        reduction_op='Sum',
        axis=1,
    )

    l = torch.arange(L).to(x)
    v, x, l = _broadcast_dims(v, x, l)
    v = _c2r(v)
    x = _c2r(x)
    l = _c2r(l)

    r = vandermonde_mult(v, x, l, backend='GPU')

    return r

L, N = 2000, 1000
v = torch.rand(N, dtype=torch.complex64)
x = 1j * torch.rand(N, dtype=torch.float32)

out = log_vandermonde(v, x, L)

print("ok")
print("shape of output : ",out.shape)
print("norm of output : ",torch.norm(out))

Also note that we have a ComplexExp1j operation in KeOps for computing $e^{ix}$ with real $x$. So you could perform the same operation faster by using the expression ComplexMult(v, ComplexExp1j(Mult(x, l))) if x and l are real-valued.

Let me know if the code above works on your side or if there is another issue.

[joanglaunes]

N.B. I just did a commit for this issue ; in fact it is a fix for the ComplexExp1j operation I was referring to in the last part of my previous comment.

[joanglaunes]

I'm closing this issue, based on my answer above. Please feel free to reopen.