如何在没有 numpy 的情况下在 jit 装饰器上设置二维数组？答案

【问题标题】：How to set 2d array on jit decorator without numpy?如何在没有 numpy 的情况下在 jit 装饰器上设置二维数组？
【发布时间】：2019-06-28 09:42:32
【问题描述】：

我在 Python3 中使用 Numba 库。

函数的参数是一个二维数组。

我将 Numba jit 装饰器设置为 list[list[int]]，但在运行代码后显示 TypeError: 'type' object is not subscriptable。

我使用print(numba.typeof(matrix))检测参数类型，它返回list(reflected list(int32))类型。

但即使我将装饰器更改为 list[list[numba.int32]] ，也无法正常工作。

代码：

from numba import jit

size = 3
matrix = [[0, 1, 2], [4, 5, 6], [7, 8, 9]]


@jit(list[list[int]])
def test(jitmatrix):
    _total = 0
    for i in range(size):
        for j in range(size):
            _total += jitmatrix[j][i]


test(matrix)

有没有想法在没有 numpy 库的 jit 装饰器上设置二维数组？

还是必须使用numpy库？

【问题讨论】：

正如我在最近的回答 (stackoverflow.com/a/56794390/392949) 中提到的，numba 不支持列表列表。在 Numba 认为二维数组的意义上，您传递的不是二维数组。如果您要传递 np.array(matrix) 并将类型规范放在装饰器中，numba 将能够解释它并 jit 代码。

标签： python arrays python-3.x numpy numba

【解决方案1】：

从 0.44 开始，Numba 不支持将列表列表作为 nopython 模式下函数的输入。见：

http://numba.pydata.org/numba-doc/latest/reference/pysupported.html#list-reflection

在@jit 的参数中，numba 不知道list 并且无法将其自动转换为任何 numba 类型。 TypeError ... subscriptable 错误来自 python 本身，因为您试图访问内置类型的元素（在本例中为 list），这是不允许的。

以下方法虽然可行：

from numba import jit
import numba as nb
import numpy as np

size = 3
matrix = np.array([[0, 1, 2], [4, 5, 6], [7, 8, 9]])


@jit(nopython=True)
# or @jit(nb.int64(nb.int64[:,:]))
def test(jitmatrix):
    _total = 0
    for i in range(size):
        for j in range(size):
            _total += jitmatrix[j,i]  # note the change in indexing, which is faster

    return _total


test(matrix)

【讨论】：

感谢您的解释和正确的示例，它让我知道我对 numba 有什么误解，让我更清楚。