How do I use numba on a member function of a class? How do I use numba on a member function of a class? python python

How do I use numba on a member function of a class?


I was in a very similar situation and I found a way to use a Numba-JITed function inside of a class.

The trick is to use a static method, since this kind of methods are not called prepending the object instance to the argument list. The downside of not having access to self is that you cannot use variables defined outside of the method. So you have to pass them to the static method from a calling method that has access to self. In my case I did not need to define a wrapper method. I just had to split the method I wanted to JIT compile into two methods.

In the case of your example, the solution would be:

from numba import jitclass MyClass:    def __init__(self):        self.k = 1    def calculation(self):        k = self.k        return self.complicated([1,2,3],k)    @staticmethod    @jit(nopython=True)                                 def complicated(x,k):                                          for a in x:            b = a**2 .+ a**3 .+ k


You have several options:

Use a jitclass (http://numba.pydata.org/numba-doc/0.30.1/user/jitclass.html) to "numba-ize" the whole thing.

Or make the member function a wrapper and pass the member variables through:

import numba as nb@nb.jitdef _complicated(x, k):    for a in x:        b = a**2.+a**3.+kclass myClass(object):    def __init__(self):        self.k = 1    def complicated(self,x):                                          _complicated(x, self.k)