numba.jit (nopython) 可以处理数组索引和切片吗？

Question

我正在尝试使用 numba.jit 加速一个简单的 Python 循环。但似乎 jit 无法处理基本的数组索引和切片？我能做些什么来使这成为可能吗？我不明白如果 jit 不能处理基本的 numpy 数组，它怎么会有用。

我正在强制使用 nopython 模式。它在对象模式下工作，但这根本不会加速代码，所以它是我需要的 nopython 模式。

下面的代码只是一个说明问题的例子。我的实际代码有同样的问题，但有更多的循环和迭代，因此 jit 会非常有用。

import numpy as np
from numba import jit

n = 100
myarray = np.zeros(n)

@jit(nopython=True)
def compute(n):
    for i in xrange(n):
        myarray[i] += 1  # This indexing causes the error.

compute(n)

"""Sample run:
> python jit_test.py
> ...
> Failed at nopython (nopython frontend)
> Internal error at <numba.typeinfer.SetItemConstrain object at 0x7f700c89a7d0>:
> Immutable array
> File "jit_test.py", line 10
"""

如果我尝试切片而不是出现不同的错误。

# The rest of the code unchanged.
        myarray[:] += 1  # This slicing causes the error.
# The rest of the code unchanged.
"""Sample run:
> python jit_test.py
> ...
> Failed at nopython (nopython frontend)
> Internal error at <numba.typeinfer.IntrinsicCallConstrain object at 0x7f297e2a9dd0>:
> Don't know how to create implicit output array with 'A' layout.
> File "jit_test2.py", line 10
"""

我的构建：

numba version: 0.19.1
numpy version: 1.9.2
python version: 2.7.10

Answer 1

在 numba 中，全局数组是静态的。您将在您的方法中获得一个只读副本。参见：

http://numba.pydata.org/numba-doc/0.20.0/reference/pysemantics.html?highlight=global#global-and-closure-variables

如果你想修改一个数组，要么显式传入它，要么在函数中创建它（在新版本的 Numba 中进行数组内存管理）和 return 它。

例如：

import numpy as np
import numba as nb

def compute(arr):
    for i in xrange(arr.shape[0]):
        arr[i] += 1

n = 100
myarray = np.zeros(n)

jitcompute = nb.jit(nopython=True)(compute)

然后计时：

In [12]: %timeit compute(myarray)
10000 loops, best of 3: 25.7 µs per loop

In [13]: %timeit jitcompute(myarray)
The slowest run took 17.06 times longer than the fastest. This could mean that an intermediate result is being cached
1000000 loops, best of 3: 461 ns per loop

一些注意事项：

这显然是一个简单的示例，但它展示了如何在方法之外定义数组，然后让函数就地修改它。
我在这里没有使用 @jit 装饰器，这样我就可以编写一次方法定义，并在有和没有 numba jit-ing 的情况下对方法进行计时，但是 jitcompute 等同于你如果您在 compute 上使用装饰器，将会得到。
如果您曾对 numba 代码进行基准测试，请确保在进行计时之前运行该方法一次，否则您将看到的是 jit 代码所花费的时间与实际执行时间处理时间。第二次运行时，您只会看到执行时间。

numba.jit (nopython) 可以处理数组索引和切片吗？

Can numba.jit (nopython) handle array indexing and slicing?

python

arrays

numpy

numba