numba jit signature of "self" for a function in class

user_16922127 0
Viewed 101

I am trying to declare the signature of a function inside a class. I am not sure what should I declare for the "self"

for example:

@jit (int32(??,int32)) def a (self, number:int) -> int

I have been searching on google without luck. should I declare the class with @jitclass decorator in order to make it work?

Thanks in advance.

numba jit python
1 Answers

As I know of, you have two solutions:

The first: declare your method as static and pass on, as arguments, every element of the class you need.

from numba import jit
arg1 = 12.1

class Thing:
    def __init__(self, arg1: np.float64):
        self.arg1 = arg1

    @staticmethod
    @jit(float64(int8, float64), nopython=True)
    def fun(number, arg1):
        # run your code
        return arg1 * number

thing = Thing(arg1=arg1)
thing.fun(number=1, arg1=thing.arg1)

The second: as you said, use a jitclass.

from numba.experimental import jitclass
arg1 = 12.1

@jitclass(spec={"arg1": float64})
class ThingJIT:
    def __init__(self, arg1):
        self.arg1 = arg1

    def fun(self, number):
        # run your code
        return self.arg1 * number

thing_jit = ThingJIT(arg1=arg1) # can use keyword arguments in jitclass __init__
thing_jit.fun(1)                # can *not* use keyword arguments in jitclass methods

There are a few drawbacks for the jitclass method:

  • You need every method to be compatible with the nopython=True parameter.
  • You can not use this class as parent or child class, even of another jitclass.
  • And I think you can not, yet, parallelise with pranges the calculations.

Hope I helped.

user_192405820
Related