How call a `@guvectorize` inside a `@guvectorize` in numba?

问题

I'm trying to call a @guvectorize inside a @guvectorize but I have an error saying :

Untyped global name 'regNL_nb': cannot determine Numba type of <class 'numpy.ufunc'>

File "di.py", line 12:
def H2Delay_nb(S1, S2, R2):
    H1 = regNL_nb(S1, S2)
    ^

here is an MRE:

import numpy as np
from numba import guvectorize, float64, int64, njit, cuda, jit

@guvectorize(["float64[:], float64[:], float64[:]"], '(n),(n)->(n)')
def regNL_nb(S1, S2, h2):
    for i in range(len(S1)):
        h2[i] = S1[i] + S2[i]

@guvectorize(["float64[:], float64[:],  float64[:]"], '(n),(n)->(n)',nopython=True)
def H2Delay_nb(S1, S2, R2):
    H1 = regNL_nb(S1, S2)
    H2 = regNL_nb(S1, S2,)
    for i in range(len(S1)):
        R2[i] =  H1[i] + H2[i]


S1 = np.array([1,2,3,4,5,6,7,8,9])
S2 = np.array([1,2,3,4,5,6,7,8,9])
H2 = H2Delay_nb(S1, S2)
print(H2)

I don't know how do I tell to numba that the function regNL_nb is a guvectorize function.

回答1:

My answer is only for the case if you're fine with replacing @guvectorize with @njit, it will be totally same code, same fast, just a bit more longer syntax to use.

It is probably some issue with accepting @guvectorize-ed functions inside other guvectorized function in nopython mode.

But Numba accepts perfectly good just regular @njit-ed functions inside other njited. So you may rewrite your function to use @njit, your function signature will remain same as @guvectorize-ed for outside world. @njit version will just need extra usage of np.empty_like(...) + return inside function.

To remind you - all @njit-ed functions are always having nopython mode enabled, so your njited code will be same fast as guvectorize+nopython.

Also I provide CUDA solution as second code snippet.

You may also make @njited only internal helper function, but external you probably can still have as @guvectorize-ed. Also if you want universal function (accepting any inputs) just remove signature 'f8[:](f8[:], f8[:])' from njited definition, signature will be auto-resolved on call.

Final code looks like this:

Try it online!

import numpy as np
from numba import guvectorize, float64, int64, njit, cuda, jit

@njit('f8[:](f8[:], f8[:])', cache = True)
def regNL_nb(S1, S2):
    h2 = np.empty_like(S1)
    for i in range(len(S1)):
        h2[i] = S1[i] + S2[i]
    return h2
        
@njit('f8[:](f8[:], f8[:])', cache = True)
def H2Delay_nb(S1, S2):
    H1 = regNL_nb(S1, S2)
    H2 = regNL_nb(S1, S2)
    R2 = np.empty_like(H1)
    for i in range(len(S1)):
        R2[i] =  H1[i] + H2[i]
    return R2

S1 = np.array([1,2,3,4,5,6,7,8,9], dtype = np.float64)
S2 = np.array([1,2,3,4,5,6,7,8,9], dtype = np.float64)
H2 = H2Delay_nb(S1, S2)
print(H2)

Output:

[ 4.  8. 12. 16. 20. 24. 28. 32. 36.]

CUDA variant of same code, it needs extra functions-wrappers if you want to automatically create and return resulting array, because CUDA-code function doesn't allow to have return value:

import numpy as np
from numba import guvectorize, float64, int64, njit, cuda, jit

@cuda.jit('void(f8[:], f8[:], f8[:])', cache = True)
def regNL_nb_cu(S1, S2, h2):
    for i in range(len(S1)):
        h2[i] = S1[i] + S2[i]
        
@njit('f8[:](f8[:], f8[:])', cache = True)
def regNL_nb(S1, S2):
    h2 = np.empty_like(S1)
    regNL_nb_cu(S1, S2, h2)
    return h2
        
@cuda.jit('void(f8[:], f8[:], f8[:])', cache = True)
def H2Delay_nb_cu(S1, S2, R2):
    H1 = regNL_nb(S1, S2)
    H2 = regNL_nb(S1, S2)
    for i in range(len(S1)):
        R2[i] =  H1[i] + H2[i]
        
@njit('f8[:](f8[:], f8[:])', cache = True)
def H2Delay_nb(S1, S2):
    R2 = np.empty_like(S1)
    H2Delay_nb_cu(S1, S2, R2)
    return R2

S1 = np.array([1,2,3,4,5,6,7,8,9], dtype = np.float64)
S2 = np.array([1,2,3,4,5,6,7,8,9], dtype = np.float64)
H2 = H2Delay_nb(S1, S2)
print(H2)

回答2:

@guvectorize(["float64[:], float64[:],  float64[:]"], '(n),(n)->(n)',nopython=True)
def H2Delay_nb(S1, S2, R2):
    H1 = regNL_nb(S1, S2)
    H2 = regNL_nb(S1, S2,)
    for i in range(len(S1)):
        R2[i] =  H1[i] + H2[i]

By using the parameter nopython = True you deactivate the object mode and hence Numba isn't able to handle all values as Python objects (refer to: https://numba.pydata.org/numba-doc/latest/glossary.html#term-object-mode)

In general are Panda, Numba or other function calls not possible if you use nopython = True. There is only a limited amount of libraries you can use with Numba Jit (in nopython). The full list can be found here: https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html.

So what you are trying to do isn't possible, except of disabling nopython, that is:

@guvectorize(["float64[:], float64[:],  float64[:]"], '(n),(n)->(n)',nopython=False)
    def H2Delay_nb(S1, S2, R2):
        H1 = regNL_nb(S1, S2)
        H2 = regNL_nb(S1, S2,)
        for i in range(len(S1)):
            R2[i] =  H1[i] + H2[i]

Going with this approach the programme outputs the correct value, i.e. [ 4. 8. 12. 16. 20. 24. 28. 32. 36.] for H2.

I also found another StackOverflow question dealing with a familiar issue: numba - TypingError: cannot determine Numba type of <class 'builtin_function_or_method'>. Credits where credit is due: Kevin K. suggested in the mentioned thread that you shall use 'simpler' data types - which are most often found in CPython. Other than that, and I fully agree with him on this point, I wouldn't be aware of any possible solution in nopython mode activated.

Sources:

numba - TypingError: cannot determine Numba type of <class 'builtin_function_or_method'>
https://numba.pydata.org/numba-doc/latest/glossary.html
https://numba.pydata.org/numba-doc/latest/proposals/type-inference.html
https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
https://github.com/numba/numba/issues/4627

来源：https://stackoverflow.com/questions/65353407/how-call-a-guvectorize-inside-a-guvectorize-in-numba

标签

python

numba