python - 索引numpy数组时numba @njit错误

Question

我正在尝试使用numba构建一个函数，该函数返回在另一个数组上求值的numpy数组
我将发布一个不带njit的简单代码:

import numpy as np
import numba as nb

def prueba(arr, eva):
    mask = []
    for i in range(len(arr)):
        mask.append(arr[i])
    return eva[mask]

如预期的那样，它可以正常工作:

>>> prueba(np.array([1,2,3]), np.array([5,6,7,8,9,10]))
array([6, 7, 8])

不过，当我尝试在nopython模式(@njit)中使用numba进行编译时，会引发错误

@nb.njit
def prueba(arr, eva):
    mask = []
    for i in range(len(arr)):
        mask.append(arr[i])
    return eva[mask]

>>> prueba(np.array([1,2,3]), np.array([5,6,7,8,9,10]))

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
<ipython-input-9-111474f08921> in <module>
----> 1 prueba(np.array([1,2,3]), np.array([5,6,7,8,9,10]))

~/.local/lib/python3.7/site-packages/numba/dispatcher.py in _compile_for_args(self, *args, **kws)
    399                 e.patch_message(msg)
    400 
--> 401             error_rewrite(e, 'typing')
    402         except errors.UnsupportedError as e:
    403             # Something unsupported is present in the user code, add help info

~/.local/lib/python3.7/site-packages/numba/dispatcher.py in error_rewrite(e, issue_type)
    342                 raise e
    343             else:
--> 344                 reraise(type(e), e, None)
    345 
    346         argtypes = []

~/.local/lib/python3.7/site-packages/numba/six.py in reraise(tp, value, tb)
    666             value = tp()
    667         if value.__traceback__ is not tb:
--> 668             raise value.with_traceback(tb)
    669         raise value
    670 

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (array(int64, 1d, C), list(int64))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
In definition 8:
    All templates rejected with literals.
In definition 9:
    All templates rejected without literals.
In definition 10:
    All templates rejected with literals.
In definition 11:
    All templates rejected without literals.
In definition 12:
    TypeError: unsupported array index type list(int64) in [list(int64)]
    raised from /home/donielix/.local/lib/python3.7/site-packages/numba/typing/arraydecl.py:71
In definition 13:
    TypeError: unsupported array index type list(int64) in [list(int64)]
    raised from /home/donielix/.local/lib/python3.7/site-packages/numba/typing/arraydecl.py:71
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of intrinsic-call at <ipython-input-8-1b5c9f1a65d5> (6)
[2] During: typing of static-get-item at <ipython-input-8-1b5c9f1a65d5> (6)

File "<ipython-input-8-1b5c9f1a65d5>", line 6:
def prueba(arr, eva):
    <source elided>
        mask.append(arr[i])
    return eva[mask]
    ^

所以我的问题是，为什么这个简单的代码会产生意外错误？我应该如何解决这个问题？

Answer 1

直接从文档中:

A subset of advanced indexing is also supported: only one advanced index is allowed, and it has to be a one-dimensional array (it can be combined with an arbitrary number of basic indices as well). https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html#array-access

因此，要使您的代码正常工作，您必须将 mask转换为 numpy array:

@nb.njit
def prueba(arr, eva):
    mask = []
    for i in range(len(arr)):
        mask.append(arr[i])
    mask_as_array = np.array(mask)
    return eva[mask_as_array]

prueba(np.array([1,2,3]), np.array([5,6,7,8,9,10]))