Flexible @generate_jit specialization on @jitclasses types?

Numba Support: How do I do ...?
1

Hi there, I’m wondering if we can @generate_jit on @jitclass types (in a Julia’s “multiple dispatch” style). The following example

@jitclass([])
class A(object):
    def __init__(self):
        pass

@jitclass([])
class B(object):
    def __init__(self):
        pass

@njit
def fa(state):
    print("Found an A")

@njit
def fb(state):
    print("Found a B")

@njit
def fc(state):
    print("Found an unknown")

@generated_jit(nopython=True)
def func(state):
    if isinstance(state, A):
        return fa
    elif isinstance(state, B):
        return fb
    else:
        return fc

func(A())
func(B())

just prints

Found an unknown
Found an unknown

I’ve tried also to use

A.class_type.instance_type

in the isinstance calls but it doesn’t help.

Cheers !

2

it seems to me this should work… when I use just the two class definitions and add

print(isinstance(A(), A))

returns True for me.
maybe try @overload?

3

Yes it works in python mode but not inside a jit function, and raises a

Untyped global name 'isinstance': cannot determine Numba type of <class 'builtin_function_or_method'>

This might explain what @generated_jit does not work with this.
If I understand well, this might come from the fact that jitclass really are experimental, and not really a “class” per se but rather implemented as a function (constructor)… But my understanding stops there

4

This works (at least in 0.57.1):

@generated_jit
def func(x):
    if isinstance(x, types.ClassInstanceType):
        if x.class_type is A.class_type:
            return fa
        elif x.class_type is B.class_type:
            return fb
        else:
            raise TypeError
    else:
        raise TypeError