I would like to create a numba jitclass with an attribute that contains any jitted function.
# simple jitted functions defined in another file
@njit
def my_function(x):
x = x + 1
return x
@njit
def another_function(x):
x = x * 2
return x
spec = [('attribute', ???),
('value', float32)]
@jitclass(spec)
class Myclass:
def __init__(self, fun):
self.attribute = fun
def class_fun(self, x):
value = self.attribute(x)
return value
When I replace ??? with numba.typeof(my_function) and create an instance of Myclass with an_object = Myclass(fun=my_function), everything works fine. But, then I can only pass the exact function my_function. When I create a Myclass object with another jitted function,
new_object = Myclass(fun=another_function)
new_object.class_fun(2)
I get the following error:
Failed in nopython mode pipeline (step: nopython mode backend)
Cannot cast type(CPUDispatcher(<function another_function at 0x7fb1e3d8ce50>))
to type(CPUDispatcher(<function my_function at 0x7fb1e404de50>))
This makes sense to me, as I defined the field type for my_function. I don't know how to define the field type attribute in a general way so that I can pass any function.
Does anyone know a way that enables to pass any jitted function when creating an object of the jitclass Myclass?