How to specify field of a numba jitclass when it contains a jitted function?

Viewed 109

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?

0 Answers
Related