Search code examples
pythonscipycudagpunumba

How to call Scipy Numba functions on GPU?

After installing the numba-scipy package, the following code snippet works:

import numba
import scipy.special as sc

@numba.vectorize(['float64(float64, float64)'])
def t_quantile(df, p):
    return sc.stdtrit(df, p)

t_quantile(2., 0.975)

How do I execute t_quantile() on the GPU, as I unsuccessfully tried in the following code?

import numba
import scipy.special as sc

@numba.vectorize(['float64(float64, float64)'], target='cuda')  # move to GPU here
def t_quantile(df, p):
    return sc.stdtrit(df, p)

t_quantile(2., 0.975)
Solution

  • scipy.special.stdtrit is actually not supported by Numba in nopython mode. Only the fallback implementation can be used (which mostly defeat the purpose of using Numba). In fact, Scipy is mostly unsupported by Numba yet. Thus, you cannot use it in Numba CUDA-targeted functions too. You could try to implement this yourself but this is pretty difficult to do.