Numba NJIT дать мой и ошибка при двумерной индексации np.array - PullRequest
0 голосов
/ 16 апреля 2019

Я пытаюсь индексировать 2D матрицу 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() декоратора

1 Ответ

1 голос
/ 16 апреля 2019

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
...