Function Pointer Inside a Jitclass

Dear all,
I am trying to create a class similar to the one below.

spec = [("x", nb.float64), ("y", nb.float64)]
@nb.experimental.jitclass(spec = spec)
class MyClass:
    def __init__(self, x, square):
        self.x = x
        if square:
            f = self.x_squared
        else:
            f = self.x_simple
        self.y = f()     

    def x_simple(self):
        return self.x
    
    def x_squared(self):
        return self.x**2

object = MyClass(x = 2.1, square = True)
print("x = ", object.x)
print("y = ", object.y)

But it raises an error when trying to compile the call of the f function. Is there a way to use a variable inside a jitclass as a function pointer to the methods of the same jitclass?

Thank you for your help,
Arthur.

hi @ArthurRocha , I haven’t tried it but it might be possible by explicitly declaring f in the spec to be a function. This thread will show you how to declare variable to be of function type: Typed list of jitted functions in jitclass - #3 by hgrecco

Hope this helps, don’t hesitate to share the result and final code! it might help the next person!

Luk