numba jit: failed type inference due to: non-precise type pyobject

I try to speed up the following function that saves the index position of the different labels present in a (large) 3 dimensional np.ndarray using jit.

import numpy as np
from numba import jit

@jit
def find_position(x, values):
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(x.shape[2]):
                values[x[i,j,k]].append((i,j))
    return values

labels = [1,2,3]
values = {l: [] for l in labels}
x = np.random.choice(labels,1000).reshape((10,10,10))

v = find_position(x, values)

However, I end up with the following error message:

Compilation is falling back to object mode WITH looplifting enabled because Function "find_position" failed type inference due to: non-precise type pyobject
During: typing of argument at <stdin> (3)

Does anyone has any tips how to get around that?

Solution

Numba cannot operate on CPython objects like lists and dicts. They are called reflected lists. This is because Numba needs to define statically the type of the objets (and they must be homogeneous) so to generate a fast native code, while CPython objects are dynamically typed and lists/dicts data structures can contain heterogeneous objects. This is also what makes CPython data structure slow compared to Numba. The two approaches are incompatible. However, CPython can deal with typed data structure at the expense of additional overheads. Numba can only operate on typed data structure that can be created from non-typed (ie. reflected) ones. The creation of such data structure is pretty slow. In the end, it may not worth it.

You can create and fill a typed dictionary like this:

import numba as nb

TupleType = nb.types.UniTuple(nb.int64, 2)
ValueType = nb.types.ListType(TupleType)
values = nb.typed.typeddict.Dict.empty(nb.int64, ValueType)
for l in labels:
    values[l] = nb.typed.typedlist.List.empty_list(TupleType)

Since values is typed, you can then pass it to Numba.

Note that dictionary are generally slow (whatever the language) as pointed out by @Rafnus. Thus, it is better to use arrays if you can, especially if the labels are small integers: you can build an array of labelMax+1 items where labelMax is the biggest label value. This assume that the label IDs are positive one. If the number of label is very small (eg. <4) and known at compile time, then the method of @Rafnus may be faster.