如何在 numba 的“@guvectorize”中调用“@guvectorize”？

Question

我正在尝试在 @guvectorize 中调用 @guvectorize，但出现错误提示：

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)
    ^

这是一个 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)

我不知道如何告诉 numba 函数 regNL_nb 是一个 guvectorize 函数。

Answer 1

我的回答仅适用于您可以将@guvectorize 替换为@njit 的情况，这将是完全相同的代码，同样的速度，只是要使用更长的语法。

在 nopython 模式下在其他 guvectorized 函数中接受 @guvectorize-ed 函数可能存在一些问题。

但是 Numba 在其他 njited 中接受非常好的常规 @njit-ed 函数。因此，您可以重写您的函数以使用@njit，您的函数签名将与外部世界的@guvectorize-ed 相同。 @njit 版本只需要额外使用np.empty_like(...) + return inside function。

提醒您 - 所有@njit-ed 函数始终启用 nopython 模式，因此您的 njit 代码将与 guvectorize+nopython 一样快。

我还提供 CUDA 解决方案作为第二个代码 sn-p。

您也可以将@njited 设为仅内部辅助函数，但您可能仍可以将外部函数设为@guvectorize-ed。此外，如果您想要通用功能（接受任何输入）只需从 njited 定义中删除签名 'f8[:](f8[:], f8[:])'，签名将在调用时自动解析。

最终代码如下所示：

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)

输出：

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

相同代码的CUDA变体，如果你想自动创建并返回结果数组，它需要额外的函数包装器，因为CUDA代码函数不允许有返回值：

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)

Answer 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]

通过使用参数nopython = True，您可以停用对象模式，因此 Numba 无法将所有值作为 Python 对象处理（请参阅：https://numba.pydata.org/numba-doc/latest/glossary.html#term-object-mode）

如果您使用nopython = True，通常无法调用 Panda、Numba 或其他函数。只有有限数量的库可以与 Numba Jit 一起使用（在 nopython 中）。 完整列表可在此处找到：https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html。

因此，除了禁用nopython 之外，您尝试执行的操作是不可能的，即：

@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]

使用这种方法，程序会输出正确的值，即 [ 4. 8. 12. 16. 20. 24. 28. 32. 36.] 对应于 H2。

我还发现了另一个处理熟悉问题的 StackOverflow 问题：numba - TypingError: cannot determine Numba type of <class 'builtin_function_or_method'>。信用到期的信用：Kevin K. 在提到的线程中建议您应该使用 'simpler' 数据类型 - 这是 CPython 中最常见的。除此之外，我完全同意他在这一点上的看法，我不知道在nopython 模式激活时有任何可能的解决方案。

---

来源：

• 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