Hey @Azmisov,
Like you I am surprised as well that some definitions fail which intuitively should work. When defining types for a Dict within a jitclass, it’s probably favorable to explicitly specify Numba types for your container members. This method is recommended by the documentation as well. The sometimes surprising behavior in your tests might be due to two reasons:
- If you don’t explicitly define field members in a jitclass, they should be Python types or another jitclass. Numba infers types in jitclasses using
as_numba_type
, but it can sometimes struggle to infer more complex datatypes, leading to unexpected outcomes.
- Method type annotations (e.g. those of init) are ignored. Within class methods Numba might not fully understand more complex datatypes within typed containers like Tuple, UniTuple, or UnicodeCharSeq. As a result, method type annotations, such as those in your examples, might not work as expected.
See further explanations below.
import numba as nb
from numba import int32
from numba.types import UnicodeCharSeq, DictType, UniTuple, Tuple
from numba.typed import Dict
from numba.experimental import jitclass
# =============================================================================
# Specifying numba.typed containers as class members explicitly
# https://numba.readthedocs.io/en/stable/user/jitclass.html#specifying-numba-typed-containers-as-class-members-explicitly
# =============================================================================
key_value_types_A = (UniTuple(int32, 2), int32)
@jitclass([('index', DictType(*key_value_types_A))])
class Test_A:
def __init__(self):
self.index = Dict.empty(*key_value_types_A)
foo = Test_A()
print(type(foo.index))
# <class 'numba.typed.typeddict.Dict'>
# =============================================================================
key_value_types_B = (Tuple((UnicodeCharSeq(6), UnicodeCharSeq(10))), int32)
@jitclass([('index', DictType(*key_value_types_B))])
class Test_B:
def __init__(self):
self.index = Dict.empty(*key_value_types_B)
foo = Test_B()
print(type(foo.index))
# <class 'numba.typed.typeddict.Dict'>
# =============================================================================
# Examples
# =============================================================================
# @jitclass
# class Test1:
# index: DictType((int32, int32), int32)
# def __init__(self):
# self.index = Dict.empty(key_type=(int32, int32), value_type=int32)
nb.extending.as_numba_type(DictType((int32, int32), int32))
# AttributeError: 'tuple' object has no attribute 'name'
# => In Test1 index must be Python type or another jitclass; field types inferred by as_numba_type fails
# https://numba.readthedocs.io/en/stable/user/jitclass.html#numba.experimental.jitclass
# =============================================================================
@jitclass
class Test2:
index: DictType(UniTuple(int32, 2), int32)
def __init__(self):
self.index = Dict.empty(key_type=UniTuple(int32, nb.int32(2)), value_type=int32)
Test2()
# TypingError: Failed in nopython mode pipeline (step: nopython frontend)
# => Method type annotations (e.g. those of __init__) are ignored.
# Not able to infer type from UniTuple(int32, nb.int32(2)) inside __init__
# https://numba.readthedocs.io/en/stable/user/jitclass.html#numba.experimental.jitclass
# =============================================================================
@jitclass
class Test3:
index: DictType(UniTuple(int32, 2), int32)
def __init__(self):
self.index = Dict.empty(key_type=(int32, int32), value_type=int32)
foo = Test3()
print(type(foo.index))
# <class 'numba.typed.typeddict.Dict'>
# =============================================================================
key_type4 = UniTuple(int32, 2)
@jitclass
class Test4:
index: DictType(key_type4, int32)
def __init__(self):
self.index = Dict.empty(key_type=key_type4, value_type=int32)
foo = Test4()
print(type(foo.index))
# <class 'numba.typed.typeddict.Dict'>
# =============================================================================
@jitclass
class Test5:
index: DictType(Tuple((UnicodeCharSeq(6), UnicodeCharSeq(10))), int32)
def __init__(self):
self.index = Dict.empty(
key_type=Tuple((UnicodeCharSeq(6), UnicodeCharSeq(10))),
value_type=int32)
Test5()
# TypingError: Failed in nopython mode pipeline (step: nopython frontend)
# => Method type annotations (e.g. those of __init__) are ignored.
# Not able to infer type from Tuple((UnicodeCharSeq(6), UnicodeCharSeq(10))) inside __init__
# https://numba.readthedocs.io/en/stable/user/jitclass.html#numba.experimental.jitclass
# =============================================================================
@jitclass
class Test6:
index: DictType(Tuple((UnicodeCharSeq(6), UnicodeCharSeq(10))), int32)
def __init__(self):
self.index = Dict.empty(
key_type=(UnicodeCharSeq(6), UnicodeCharSeq(10)),
value_type=int32)
Test6()
# TypingError: Failed in nopython mode pipeline (step: nopython frontend)
# => Method type annotations (e.g. those of __init__) are ignored.
# Not able to infer type from (UnicodeCharSeq(6), UnicodeCharSeq(10)) inside __init__
# https://numba.readthedocs.io/en/stable/user/jitclass.html#numba.experimental.jitclass
# =============================================================================
# works
key_type7 = Tuple((UnicodeCharSeq(6), UnicodeCharSeq(10)))
@jitclass
class Test7:
index: DictType(key_type7, int32)
def __init__(self):
self.index = Dict.empty(key_type=key_type7, value_type=int32)
foo = Test7()
print(type(foo.index))
# <class 'numba.typed.typeddict.Dict'>