55

I'm using the stable version of Numba 0.30.1.

I can do this:

import numba as nb
@nb.jit("void(f8[:])",nopython=True)                             
def complicated(x):                                  
    for a in x:
        b = a**2.+a**3.

as a test case, and the speedup is enormous. But I don't know how to proceed if I need to speed up a function inside a class. 

import numba as nb
def myClass(object):
    def __init__(self):
        self.k = 1
    #@nb.jit(???,nopython=True)                             
    def complicated(self,x):                                  
        for a in x:
            b = a**2.+a**3.+self.k

What numba type do I use for the self object? I need to have this function inside a class since it needs to access a member variable.

dbrane
  • 887
  • 1
  • 7
  • 19
  • What about a [`jitclass`](http://numba.pydata.org/numba-doc/dev/user/jitclass.html)? I don't think it's otherwise possible to avoid the "object-fallback" given that `self` is by definition a _object_. – MSeifert Jan 20 '17 at 17:38
  • Also what is `b = a**2.+a**3.+self.k` going to achieve given that you immediatly overwrite it in the next loop? – MSeifert Jan 20 '17 at 17:47
  • 2
    The `self.k` is just to schematically show that I will need to call member variables, and can't just have the function outside the class – dbrane Jan 20 '17 at 23:29

2 Answers2

62

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 jit

class 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
Marduk
  • 982
  • 8
  • 12
31

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.jit
def _complicated(x, k):
    for a in x:
        b = a**2.+a**3.+k

class myClass(object):
    def __init__(self):
        self.k = 1

    def complicated(self,x):                                  
        _complicated(x, self.k)
JoshAdel
  • 66,734
  • 27
  • 141
  • 140
  • 8
    I've seen that page on `jitclass` but it's completely unclear there how I can explicitly state the data types of each member function. Could you show an example? The wrapper function approach becomes inelegant and beats the point of me having put things in a class to begin with. – dbrane Jan 20 '17 at 23:31
  • I would avoid declaring the types of the arguments to member function and just Numba handle it via type inference. In recent memory, I can't recall declaring types to result in better performance. There are some more complex jitclass examples in https://github.com/numba/numba/tree/master/examples -- e.g. https://github.com/numba/numba/blob/master/examples/stack.py – JoshAdel Jan 21 '17 at 02:25
  • @dbrane as I know the methods of the class can not be jitted in current version of numba. If try to do it pretty clear error message appears "TypeError: class members are not yet supported: complicated". Also, I've found that usage of jitclass, does not give any speedup for your example. – Roman Fursenko Feb 01 '17 at 07:39
  • I declared the helper function inside the member function of the class I am using: `def complicated(self, x): def calc(x, k): return some difficult calculation` – Sebastiano1991 Apr 17 '19 at 08:37
  • Late to the party but sometimes you have to specify the variable types and the type definitions that numba needs for that are [here](https://github.com/numba/numba/blob/61ec1fd0f69aeadece218dccf4c39ebc5c7dfbc4/docs/source/reference/types.rst). Not sure why numba dynamically infers the types when applied to a plain function but needs those type definitions in the case of a class. – matanster Jun 14 '22 at 20:11