使用 Numba 提取 numpy 数组中的特定行

Question

I have the following array:

import numpy as np
from numba import njit


test_array = np.random.rand(4, 10)

I create a "jitted" function that slices the array and does some operations afterwards:

@njit(fastmath = True)
def test_function(array):

   test_array_sliced = test_array[[0,1,3]]

   return test_array_sliced

However, Numba throws the following error:

In definition 11:
    TypeError: unsupported array index type list(int64) in [list(int64)]
    raised from /Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/numba/typing/arraydecl.py:71
This error is usually caused by passing an argument of a type that is unsupported by the named function.

Workaround

I have tried to delete the rows I do not need by using np.delete , but since I have to specify an axis Numba throws the following error:

@njit(fastmath = True)
def test_function(array):

   test_array_sliced = np.delete(test_array, obj = 2, axis = 0)

   return test_array_sliced

In definition 1:
    TypeError: np_delete() got an unexpected keyword argument 'axis'
    raised from /Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/numba/typing/templates.py:475
This error is usually caused by passing an argument of a type that is unsupported by the named function.

Any ideas of how to extract specific rows under Numba?

Answer 1

I think it will work (it seems to suggest so in the docs ) if you index with an array instead of a list:

test_array_sliced = array[np.array([0,1,3])]

(I changed the array you're slicing to array , which is what you pass in to the function. Maybe it was intentional, but be careful with globals!)

Answer 2

Numba does not support numpy fancy indexing. I'm not 100% sure what your real use case looks like, but a simple way to do it would be something like:

import numpy as np
import numba as nb

@nb.njit
def test_func(x):
    idx = (0, 1, 3)
    res = np.empty((len(idx), x.shape[1]), dtype=x.dtype)
    for i, ix in enumerate(idx):
        res[i] = x[ix]

    return res

test_array = np.random.rand(4, 10)
print(test_array)
print()
print(test_func(test_array))

Edit: @kwinkunks is correct, and my original answer made an incorrect blanket statement that fancy indexing was not supported. It is in a limited set of cases, including this one.