How to use numba njit on dictionaried with list of arrays as values-CodePudding

I am trying to learn how to use python dictionaries with numba njit. In this regard, we must convert the pythonic dictionary items to numba compatible one. I got stuck in this step when working on a dictionary in which values were list of arrays. The following example can reproduce the problem:

sample_dict = {'size_0': [np.array([0.021, 0.022]), np.array([0.03, 0.008])],
               'size_1': [np.array([0.031]), np.array([0.036, 0.003])],
               'size_2': [np.array([], dtype=np.float64), np.array([0.043])]}
num_ = 2



String = nb.types.unicode_type
ValueArray = nb.float64[::1]

ValueList = nb.types.ListType(ValueArray)
ValueDict = nb.types.DictType(String, ValueList)

# DictValue = nb.types.Tuple([ValueArray])
# ValueDict = nb.types.DictType(String, DictValue)

# DictValue = nb.types.Tuple([nb.types.List(ValueArray, reflected=True)])
# ValueDict = nb.types.DictType(String, DictValue)

sample_nb = nb.typed.typeddict.Dict.empty(String, ValueArray)
for key, value in sample_dict.items():
    sample_nb[key] = value.copy()
    # sample_nb[key] = nb.types.List(value)


@nb.njit(ValueDict, nb.int_)
def nbTest(sample_dict, num_):
    ii = 1
    for i in range(num_):
        for j in sample_dict.values():
            ii  = len(j[i])


nbTest(sample_nb, num_)

I've get different errors since I have tried various ways, but the most happened one was:

No implementation of function Function() found for signature:

setitem(DictType[unicode_type,array(float64, 1d, C)]<iv=None>, unicode_type, reflected list(array(float64, 1d, C))<iv=None>)

How to prepare such dictionaries for using by numba njit? Is there any limitation if one of dictionary values have just one array (not as this example which all dictionary values contain 2 arrays)?

CodePudding user response：

The issue comes from two point: the type is not valid in the dict as it is ValueArray while it should be ValueList and the reflected list that cannot be copied into the dict using a basic assignment. For the second point nb.typed.typedlist.List can be used to fix that. Note that nb.types.List is a type while nb.typed.typedlist.List is used to create instances. Here is the resulting code:

String = nb.types.unicode_type
ValueArray = nb.float64[::1]

ValueList = nb.types.ListType(ValueArray)
ValueDict = nb.types.DictType(String, ValueList)

# DictValue = nb.types.Tuple([ValueArray])
# ValueDict = nb.types.DictType(String, DictValue)

# DictValue = nb.types.Tuple([nb.types.List(ValueArray, reflected=True)])
# ValueDict = nb.types.DictType(String, DictValue)

sample_nb = nb.typed.typeddict.Dict.empty(String, ValueList)
for key, value in sample_dict.items():
    sample_nb[key] = nb.typed.typedlist.List(value.copy())
    # sample_nb[key] = nb.types.List(value)