python - 具有复杂 numpy 数组和 native 数据类型的 numba TypingError

Question

我有一个处理复杂数据类型的函数，我正在使用 numba为了更快的处理。我使用 numpy 声明了一个零数组, 具有复杂数据类型，稍后在函数中填写。但是在运行时numba不能使零生成函数过载。为了重现错误，我提供了一个 MWE。

import numpy as np
from numba import njit

@njit
def my_func(idx):
    a = np.zeros((10, 5), dtype=complex)
    a[idx] = 10
    return a

my_func(4)

显示以下错误，其中数组 a正在初始化。

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)

No implementation of function Function(<built-in function zeros>) found for signature:
zeros(Tuple(Literal[int](10), Literal[int](5)), dtype=Function(<class 'complex'>))
There are 2 candidate implementations:

 Of which 2 did not match due to:
  Overload of function 'zeros': File: numba\core\typing\npydecl.py: Line 511.
    With argument(s): '(UniTuple(int64 x 2), dtype=Function(<class 'complex'>))':
   No match.

我假设这与变量 a 的数据类型有关。 (我需要它很复杂)。我该如何解决这个错误？
任何帮助将不胜感激，谢谢。

Answer 1

你的问题与复数无关。如果您指定 a = np.zeros((10, 5), dtype=int) ，你会遇到同样的问题。
虽然 numpy采用 python 原生数据类型 int , float和 complex并将它们视为 np.int32 , np.float64和 np.complex128 , numba然而，它本身并不这样做。
因此，每当您在 jitted 函数中指定数据类型时，您要么使用 numpy数据类型:

import numpy as np
from numba import njit

@njit
def my_func(idx):
    a = np.zeros((10, 5), dtype=np.complex128)
    a[idx] = 10
    return a

my_func(4)

或者你用 numba通过直接导入的数据类型:

import numpy as np
from numba import njit, complex128

@njit
def my_func(idx):
    a = np.zeros((10, 5), dtype=complex128)
    a[idx] = 10
    return a

my_func(4)

或通过 types :

import numpy as np
from numba import njit, types

@njit
def my_func(idx):
    a = np.zeros((10, 5), dtype=types.complex128)
    a[idx] = 10
    return a

my_func(4)

据我所知，您使用这些选项中的哪一个并没有什么区别。 Here是 numba 文档的相关部分。