Numba 编译函数无法迭代具有不同维度的 numpy 数组的元组

Numba compiled function cannot iterate over tuple of numpy arrays with different dimensions

我正在尝试修改一些具有不同维度的 numpy 数组,并使用元组传递给 numba 函数。该函数如下所示:

from numba import njit

def numbatest(tuple1,tuple2):
    for i in range(2):
        tuple1[i][0]=tuple2[i][2]
    return tuple1,tuple2

numbatest_njit=njit(numbatest)

当我将两个 numpy 数组(一维和二维)的元组传递给它时:

a=np.empty(10)
b=np.empty([10,3])
c=np.empty(10)
d=np.empty([10,3])
A=(a,b)
B=(c,d)
C,D=numbatest_njit(A,B)

我收到以下错误

TypingError                               Traceback (most recent call last)
<ipython-input-179-f6c1f66607ba> in <module>
      6 A=(a,b)
      7 B=(c,d)
----> 8 C,D=numbatest_njit(A,B)
      9 

~/.local/lib/python3.6/site-packages/numba/core/dispatcher.py in _compile_for_args(self, *args, **kws)
    418                 e.patch_message(msg)
    419 
--> 420             error_rewrite(e, 'typing')
    421         except errors.UnsupportedError as e:
    422             # Something unsupported is present in the user code, add help info

~/.local/lib/python3.6/site-packages/numba/core/dispatcher.py in error_rewrite(e, issue_type)
    359                 raise e
    360             else:
--> 361                 raise e.with_traceback(None)
    362 
    363         argtypes = []

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(Tuple(array(float64, 1d, C), array(float64, 2d, C)), int64)
 
There are 22 candidate implementations:
      - Of which 22 did not match due to:
      Overload of function 'getitem': File: <numerous>: Line N/A.
        With argument(s): '(Tuple(array(float64, 1d, C), array(float64, 2d, C)), int64)':
       No match.

During: typing of intrinsic-call at <ipython-input-178-e1c23213087e> (6)

File "<ipython-input-178-e1c23213087e>", line 6:
def numbatest(tuple1,tuple2):
    <source elided>
    for i in range(2):
        tuple1[i][0]=tuple2[i][2]
        ^

有趣的是,如果我重写该函数以便我显式调用不同的元组,该函数将起作用:

def numbatest(tuple1,tuple2):
    i=0
    tuple1[i][0]=tuple2[i][2]
    i=1
    tuple1[i][0]=tuple2[i][2]
    return tuple1,tuple2

如果直接调用第一个函数(不使用 numba 编译它),它就可以工作。如果所有的 numpy 数组都是一维数组,那么它也适用于 numba。

有谁知道这可能是什么问题?以及如何避免呢?您如何将不同大小的 numpy 数组传递给函数呢?非常感谢您!

Numba 支持 homogeneous and heterogeneous tuples 的不同操作集。特别是,当涉及到元组索引时,numba 仅支持索引 异构 元组,前提是索引是 在 compile-time 已知的常量。

这解释了元组包含一维数组(同构元组)与混合一维和二维数组(异构元组)时的行为差异。它还解释了为什么第二个版本的函数有效:numba 可以分辨出 i = 0i = 1 是 compile-time 常量,但不够聪明,无法知道 range(2) 可以展开同样的事情。

Numba 元组与 Python 的元组相对不同,因为静态类型(与 Python 中的动态类型相反)。事实上,元组中项目的类型可以不同(这实际上是你的情况)因此你不能像在 Python 中那样循环所有项目:没有类型可以匹配所有不同的项目类型。更具体地说,变量 tuple1[i] 的类型不明确,因为它应该是 array(float64, 1d, C)array(float64, 2d, C).

请注意,此类问题通常在 statically-typed 编译语言中使用变体类型来解决,但使用变体类型是不安全的,并且不会导致速度提升,因为静态类型主要是使 Numba 比CPython.