Нарезка массивов Numpy в функциях Numba

Я хочу нарезать 2D-массив Numpy внутри функции, украшенной Numba jit. Индексы строк и столбцов доступны в виде списков / массивов numpy. Вот MWE

import numpy as np
from numba import jit

@jit(nopython=True)
def func(a,rowinds,colinds):
    return np.sum(a[rowinds,colinds])


A = np.random.random([10,10])
rowinds = np.array([0,1,2])
colinds = np.array([3,4,5])

res = func(A,rowinds,colinds)

Однако это вызывает следующую ошибку:

Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(array(float64, 2d, C), UniTuple(array(int64, 1d, C) x 2))

Есть ли способ обойти это без ручного перебора элементов rowinds или colinds?


person lovespeed    schedule 16.07.2021    source источник
comment
вы хотите разрезать или хотите суммировать? если вы воспроизводите эти функции только с помощью numpy, вы в основном суммируете 3 элемента A, это то, что вы хотите?   -  person Bruno Magacho    schedule 17.07.2021
comment
@BrunoMagacho: Я хочу суммировать это, но в моем реальном коде массив имеет гораздо больший размер, а количество элементов, которые я хочу суммировать, намного больше. (пример с 3 элементами A - это просто MWE, чтобы проиллюстрировать проблему).   -  person lovespeed    schedule 18.07.2021
comment
на самом деле я не знаю, как эти типы могут быть помещены в качестве индексов, я думаю, что это средство numpy, которое не разрешено в numba, это связано с тем, что описано в [этом сообщении] (stackoverflow.com/questions/ 32508628 /). В вашем приложении есть проблема в loopify по rowinds и colinds? может быть еще более эффективным для дальнейшей оптимизации   -  person Bruno Magacho    schedule 19.07.2021