Numba – Accessing Numpy Array with variable index

numba is being really clever here! Consider what happens when you pass a col to get_index that isn’t in cols. cols[i] == col will never be true, the loop will exit, and since there’s no catchall return at the end of the function, the return value will be None.

numba therefore correctly infers that the return type of get_index is OptionalType(int64) i.e. a value that may either be int64 or None. But None isn’t a valid type for indices, so you can’t use a value that might be None to index an array.

You can fix this by adding a catchall return at the end.

@jit(nopython=True)
def get_index(cols, col):
    for i in range(len(cols)):
        if cols[i] == col:
            return i
    return -1

Of course this might not be the behavior you want in this case; it’s probably better to raise an exception, which numba also handles correctly.

@jit(nopython=True)
def get_index(cols, col):
    for i in range(len(cols)):
        if cols[i] == col:
            return i
    raise IndexError('list index out of range')

CLICK HERE to find out more related problems solutions.

Leave a Comment

Your email address will not be published.

Scroll to Top