Numba how to use dict in a class

Question

as specified below, the dict should have keys 2-tuples of int and values ints.

from numba.experimental import jitclass
import numba

@jitclass({'shape': numba.types.Tuple((numba.int32, numba.int32)), 'dict': numba.types.DictType(numba.types.UniTuple(numba.types.int32, 2), numba.int32)})
class BigramCounts:
    def __init__(self, shape: tuple[int, int]):
        self.shape = shape
        self.dict = {}  # this does not work
        
        # the following does not work either
        # self.dict = numba.typed.Dict.empty(key_type=numba.types.UniTuple(numba.types.int32, 2), value_type=numba.int32)

b_c = BigramCounts((2, 3))

Unfortunately plain self.dict = {} initialization does not work:

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Cannot infer the type of variable '$8build_map.2' (temporary variable), have imprecise type: DictType[undefined,undefined]<iv={}>. 

File "scratch_3.py", line 8:
    def __init__(self, shape: tuple[int, int]):
        <source elided>
        self.shape = shape
        self.dict = {}  # this does not work
        ^

During: resolving callee type: jitclass.BigramCounts#1055192d0<shape:UniTuple(int32 x 2),dict:DictType[UniTuple(int32 x 2),int32]<iv=None>>
During: typing of call at <string> (3)

During: resolving callee type: jitclass.BigramCounts#1055192d0<shape:UniTuple(int32 x 2),dict:DictType[UniTuple(int32 x 2),int32]<iv=None>>
During: typing of call at <string> (3)


File "<string>", line 3:
<source missing, REPL/exec in use?>

The second initialization does not work either (how come numba.types does not support classes?):

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend) Failed in nopython mode pipeline (step: nopython frontend) Invalid use of <class 'numba.core.types.containers.UniTuple'> with parameters (class(int32), Literal[int](2)) No type info available for <class 'numba.core.types.containers.UniTuple'> as a callable. During: resolving callee type: typeref[<class 'numba.core.types.containers.UniTuple'>] During: typing of call at /Users/adam/Library/Application Support/JetBrains/PyCharm2022.3/scratches/scratch_3.py (11)


File "scratch_3.py", line 11:
    def __init__(self, shape: tuple[int, int]):
        <source elided>
        # the following does not work either
        self.dict = numba.typed.Dict.empty(key_type=numba.types.UniTuple(numba.types.int32, 2), value_type=numba.int32)
        ^

During: resolving callee type: jitclass.BigramCounts#11270d5a0<shape:UniTuple(int32 x 2),dict:DictType[UniTuple(int32 x 2),int32]<iv=None>> During: typing of call at <string> (3)

During: resolving callee type: jitclass.BigramCounts#11270d5a0<shape:UniTuple(int32 x 2),dict:DictType[UniTuple(int32 x 2),int32]<iv=None>> During: typing of call at <string> (3)


File "<string>", line 3: <source missing, REPL/exec in use?>

Answer 1

You need to define the types outside the jitclass.

from numba import int32
from numba.types import Tuple, DictType, UniTuple

key_type = UniTuple(numba.types.int32, 2)
value_type = int32

@jitclass({
    'shape': Tuple((int32, int32)), 
    'dict': DictType(key_type, value_type),
})
class BigramCounts:
    def __init__(self, shape: tuple[int, int]):
        self.shape = shape
        self.dict = numba.typed.Dict.empty(
            key_type=key_type, 
            value_type=value_type,
        )

b_c = BigramCounts((2, 3))
b_c.dict[(6,7)] = 8

Or alternatively use:

self.dict = {(int32(0), int32(0)): int32(0) for _ in range(0)}