I know numba has a very simple approach to OOP by permitting structs and methods to operate on this structs. However, I have many cases where I would like to reuse code which manipulates two related structs in the same way.
Is there a way how I (1.) can define an overload on a numba jitclass function by using (2.) a shared jit implementation which is then compiled for each struct type. I’m thinking of something like this:
from numba.core.extending import intrinsic, overload, register_jitable
from numba.experimental import jitclass
@jitclass
class ObjectA:
def __init__(self, a):
self.a = a * 1
def get_a(self):
raise NotImplementedError
@jitclass
class ObjectB:
def __init__(self, a):
super().__init__(a)
self.a = a * 2
def get_a(self):
raise NotImplementedError
@register_jitable
def get_a(obj):
return obj.a
@overload(ObjectA.get_a)
def get_a_A(obj):
return get_a
@overload(ObjectB.get_a)
def get_a_B(obj):
return get_a
This should ideally result in the following behaviour:
obja = ObjectA(10)
objb = ObjectB(10)
print(obja.get_a())
# 10
print(objb.get_a())
# 20
I would also be interested if there are any other ways how one could share implementations or across structs while minimizing duplications.