Extracting specific rows in numpy array using Numba

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?

Matt Hall · Accepted Answer · 2019-10-27 15:16:01Z

2

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!)

answered Oct 27, 2019 at 15:16

Matt Hall

8,1791 gold badge28 silver badges38 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

JoshAdel · Accepted Answer · 2019-10-27 20:47:21Z

1

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

edited Oct 27, 2019 at 20:47

answered Oct 27, 2019 at 15:11

JoshAdel

69.1k27 gold badges146 silver badges146 bronze badges

Collectives™ on Stack Overflow

Extracting specific rows in numpy array using Numba

Workaround

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

Workaround

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related