Numba 不理解矢量化

Question

我刚刚开始使用 numba 来加速我的数值模拟。在模拟中，我必须对 2x2 矩阵进行矢量化，但这不能仅使用 numpy 方法来完成。请参阅下面的一个非常简单的示例：

import numpy as np
from numba import jit

def example(matrix):
    return matrix.T.reshape((4,1))

matrix = np.array([[1,2],[3,4]])

example(matrix)

这给出了我的预期结果，矩阵的矢量化形式：

array([[1],
       [3],
       [2],
       [4]])

但是，当我尝试在 'jit' 的帮助下做同样的事情时，我得到以下信息：

@jit(nopython=True)
def example(matrix):
    return matrix.T.reshape((4,1))

example(matrix)

NotImplementedError                       
Traceback (most recent call last)
<ipython-input-6-b23f267f5d22> in <module>
----> 1 example(matrix)

NotImplementedError: incompatible shape for array

我对此输出感到非常困惑，因为我只使用了 numpy，而 numba 的文档指出它同时支持“.T”和“.reshape”。一个小警告是我的矩阵存储了 numpy.complex128 值，但对于这个最小代码我也得到了同样的错误。

Answer 1

添加 copy 使其工作。在 numpy reshape after transpose 中生成一个副本，而不是一个视图，因为底层数据缓冲区被重新排序。显然 numba 没有为 reshape 添加这种灵活性。

In [161]: @numba.njit
     ...: def foo(arr):
     ...:     return arr.T.copy().reshape(-1,1)
In [162]: 
In [162]: foo(np.arange(4).reshape(2,2))
Out[162]: 
array([[0],
       [2],
       [1],
       [3]])
In [163]: timeit foo(np.arange(4).reshape(2,2))
3.52 µs ± 31.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

将此与常规 numpy 版本进行比较：

In [164]: np.arange(4).reshape(2,2).T.reshape(-1,1)
Out[164]: 
array([[0],
       [2],
       [1],
       [3]])
In [165]: timeit np.arange(4).reshape(2,2).T.reshape(-1,1)
3.18 µs ± 15.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

虽然此操作生成了 MATLAB 样式的列向量，但它并不是我们通常所说的 numpy 向量化。 [164] 已经优化使用了 numpy 编译方法。

Numba 不理解矢量化

Numba does not understand vectorization

python

numpy

matrix

numba