CODEWITHSUNDEEP
pythonnumbajit
zezo
zezo

Reputation: 455

compile python class with numba

I am using the numba.jitclass() decorator to compile a python class Honda that takes another python class Car as input argument:

class Car(object):
    cc = '10'

class Honda(object):
    def __init__(self):
        self.car = Car()

    def cal(self):
        'Do heavy calculations'

Numba can't identity the type of the type of python object (class) !

object = Car()
numba.typeof(object)      # Error (not numba type)

Trial _ 1 (Not working)

from numba.experimental import jitclass
from numba import types, jit, njit

@jitclass([('car', types.pyobject)])  # I tried to add `nopython=False`, but raise error
class Honda(object):
    def __init__(self):
        self.car = Car()
    
    @jit
    def cal(self):  # TypeError: class members are not yet supported
        'Do heavy calculations'

Trial _ 2 (Not working)

@jitclass([('car', types.pyobject)])  # I tried to add `nopython=False`, but raise error
class Honda(object):
    def __init__(self):
        self.car = typed.pyobject(Car())

    @jit
    def cal(self):  # TypeError: class members are not yet supported
        'Do heavy calculations'

Trial _ 3 (Not working)

@jitclass() 
class Honda(object, Car):
    def __init__(self):
        pass

    @jit
    def cal(self):  # TypeError: class members are not yet supported
        'Do heavy calculations'

How can I simply use any python object (of any class type) inside a class that is compiled by a numba.jitclass() ? and add class members (@njit methods) to be used ?

Is there any better way to achieve this rather than a numba.jitclass() ?

Upvotes: 1

Views: 1380

Answers (1)

aerobiomat
aerobiomat

Reputation: 3437

If you want to use a class in a jitclass, you need to make it a jitclass as well.

However:

@nb.experimental.jitclass
class Car0:
    cc = "10"

produces

TypeError: class members are not yet supported: cc

If you can live with an instance member instead of a class member, it will compile. For example:

@nb.experimental.jitclass
class Car:
    _cc: str

    def __init__(self, cc: str):
        self._cc = cc

    def cc(self) -> str:
        return self._cc + " cc"

Check:

>>> c = Car("10")
>>> c.cc()
'10 cc'

Then you can use it in the second class. Declaring its type is ugly, but it works:

car_type = nb.typeof(Car("10"))

@nb.experimental.jitclass
class Honda:
    car: car_type

    def __init__(self):
        self.car = Car("20")

    def cal(self) -> str:
        return self.car.cc()

Check:

>>> h = Honda()
>>> h.cal()
'20 cc'

Upvotes: 2

Related Questions