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

Question

我正在尝试B使用包含我想要的索引的向量来索引njit 函数中的二维矩阵， 这里a的矩阵切片是D一个最小的示例：import numba as nbimport numpy as [e-mail]()def test(N,P,B,D):    for i in range(N):        a = D[i,:]        b =  B[i,a]        P[:,i] =bP = np.zeros((5,5))B = np.random.random((5,5))*100D = (np.random.random((5,5))*5).astype(np.int32)print(D)N = 5print(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 bythe use of unsupported features or an issue in resolving types.我不明白我在这里做错了什么。代码在没有@nb.njit()装饰器的情况下工作

Answer 1

numba 不支持所有与 numpy 相同的“花式索引” - 在这种情况下，问题是使用a数组选择数组元素。

对于您的特定情况，因为您b事先知道形状，您可以像这样解决：

import numba as nb

import numpy as np

@nb.njit

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

    b = np.empty(D.shape[1], dtype=B.dtype)

    for i in range(N):

        a = D[i,:]

        for j in range(a.shape[0]):

            b[j] = B[i, j]

        P[:, i] = b

Answer 2

另一种解决方案是在调用测试之前在 B 上应用交换轴并反转索引（B[i,a]-> B[a,i]）。我不知道为什么这是有效的，但这里是实现：

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[a,i]

        P[:, i] = b

    

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

B = np.arange(25).reshape((5,5))

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

N = 5

test(N,P,np.swapaxes(B, 0, 1), D)

顺便说一句，在@chrisb 给出的答案中，不是：b[j] = B[i, j]而是b[j] = B[i, a[j]]。