numba njit模式下创建complex类型numpy零数组报TypingError如何解决

问题原因 #

numba的nopython模式下，无法将Python内置的complex类型识别为numpy数组的合法dtype参数，所以找不到对应np.zeros的重载实现。

解决方法 #

将dtype参数替换为numpy原生的复数类型即可，常用的有两种：

• 双精度复数：np.complex128（对应Python原生complex的精度，最常用）
• 单精度复数：np.complex64（内存占用更低，适合对精度要求不高的场景）

修改后的可运行代码如下：

import numpy as np
from numba import njit

@njit
def my_func(idx):
    # 替换dtype为np.complex128即可
    a = np.zeros((10, 5), dtype=np.complex128)
    a[idx] = 10
    return a

my_func(4)

补充说明 #

不光是复数类型，在numba的nopython模式下调用numpy数组创建函数时，尽量都使用numpy提供的内置类型作为dtype参数，避免使用Python原生的int/float/complex类，能大概率避免这类类型匹配错误。

内容的提问来源于stack exchange，提问作者learner