numba njit 在 2D np.array 索引上给出我的和错误

Question

我正在尝试B使用包含我想要的索引的向量来索引njit 函数中的二维矩阵， 这里a的矩阵切片是D一个最小的示例：

import numba as nb

import numpy as np

@nb.njit()

def test(N,P,B,D):

    for i in range(N):

        a = D[i,:]

        b =  B[i,a]

        P[:,i] =b

P = np.zeros((5,5))

B = np.random.random((5,5))*100

D = (np.random.random((5,5))*5).astype(np.int32)

print(D)

N = 5

print(P)

test(N,P,B,D)

print(P)

我在线路上收到 numba 错误 b =  B[i,a]

File "dj.py", line 10:

def test(N,P,B,D):

    <source elided>

        a = D[i,:]

        b =  B[i,a]

        ^

This is not usually a problem with Numba itself but instead often caused by

the use of unsupported features or an issue in resolving types.

我不明白我在这里做错了什么。代码在没有@nb.njit()装饰器的情况下工作

Answer 1

numba 不支持所有与 numpy 相同的“花式索引” - 在这种情况下，问题是使用a数组选择数组元素。对于您的特定情况，因为您b事先知道形状，您可以像这样解决：import numba as nbimport numpy as [e-mail] test(N,P,B,D):&nbsp; &nbsp; b = np.empty(D.shape[1], dtype=B.dtype)&nbsp; &nbsp; for i in range(N):&nbsp; &nbsp; &nbsp; &nbsp; a = D[i,:]&nbsp; &nbsp; &nbsp; &nbsp; for j in range(a.shape[0]):&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; b[j] = B[i, j]&nbsp; &nbsp; &nbsp; &nbsp; P[:, i] = b

Answer 2

另一种解决方案是在调用测试之前在 B 上应用交换轴并反转索引（B[i,a]-> B[a,i]）。我不知道为什么这是有效的，但这里是实现：import numba as nbimport numpy as [e-mail]()def test(N,P,B,D):&nbsp; &nbsp; for i in range(N):&nbsp; &nbsp; &nbsp; &nbsp; a = D[i,:]&nbsp; &nbsp; &nbsp; &nbsp; b =&nbsp; B[a,i]&nbsp; &nbsp; &nbsp; &nbsp; P[:, i] = b&nbsp; &nbsp;&nbsp;P = np.zeros((5,5))B = np.arange(25).reshape((5,5))D = (np.random.random((5,5))*5).astype(np.int32)N = 5test(N,P,np.swapaxes(B, 0, 1), D)顺便说一句，在@chrisb 给出的答案中，不是：b[j] = B[i, j]而是b[j] = B[i, a[j]]。