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

Question

我正在尝试修改一些具有不同维度的 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 数组传递给函数呢？非常感谢您！

Answer 1

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

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

Answer 2

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

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

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

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

python

jit

tuples

numba

numpy-ndarray