Create a Numba typed dictionary with integers as keys and arrays of float64 as values-CodePudding

I need to define a dictionary with integers as keys and arrays of float64 as values. In Python I can define it with:

import numpy as np

d = {3: np.array([0, 1, 2, 3, 4])}

To create the same type of dictionary in a Numba-compiled function I do

import numba

@numba.njit()
def generate_d():

    d = Dict.empty(types.int64, types.float64[:])

    return d

but I get an error at compile time. I don't understand why it errors, given the very simple instructions.

This is the error when I run generate_d():

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
/tmp/ipykernel_536115/3907784652.py in <module>
----> 1 generate_d()

~/envs/oasis/lib/python3.8/site-packages/numba/core/dispatcher.py in _compile_for_args(self, *args, **kws)
    466                 e.patch_message(msg)
    467 
--> 468             error_rewrite(e, 'typing')
    469         except errors.UnsupportedError as e:
    470             # Something unsupported is present in the user code, add help info

~/envs/oasis/lib/python3.8/site-packages/numba/core/dispatcher.py in error_rewrite(e, issue_type)
    407                 raise e
    408             else:
--> 409                 raise e.with_traceback(None)
    410 
    411         argtypes = []

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(class(float64), slice<a:b>)
 
There are 22 candidate implementations:
      - Of which 22 did not match due to:
      Overload of function 'getitem': File: <numerous>: Line N/A.
        With argument(s): '(class(float64), slice<a:b>)':
       No match.

During: typing of intrinsic-call at /tmp/ipykernel_536115/3046996983.py (4)
During: typing of static-get-item at /tmp/ipykernel_536115/3046996983.py (4)

File "../../../../tmp/ipykernel_536115/3046996983.py", line 4:
<source missing, REPL/exec in use?>

I get the same error even if I explicit the signature

@numba.njit("float64[:]()")
def generate_d():

    d = Dict.empty(types.int64, types.float64[:])

    return d

I'm using numba v 0.55.1, numpy 1.20.3

How can I get this to work?

CodePudding user response：

As far as I know type expressions are not supported in JIT functions yet (Numba version 0.54.1). You need to create the type outside the function. Here is an example:

import numba
from numba.typed import Dict

# Type defined outside the JIT function
FloatArrayType = numba.types.float64[:]

@numba.njit
def generate_d():
    d = Dict.empty(numba.types.int64, FloatArrayType)  # <-- and used here
    return d