Passing a class object as a function argument, in Numba optimized Python

Question

I want to pass a class object to a function. I can make it work, but I am wondering if there is a type I can assign it? I have a "minimal" example of what I am trying to do.

spec = [("a", float64),("b",float64)]
@jitclass(spec)
class SOMETHING_3():
    def __init__(self):
        self.a = 1.1
        self.b = 2.3

    def sum(self):
        return self.a + self.b


@jit(float64(float64, XXX), nopython = True)
def get_sum_3(c, someobj):
    d = 0
    for i in range(1000):
        for j in range(1000):
            d += c + someobj.sum()
    return d   

If I remove the explict type assigment "float64(float64, XXX)" it works fine.

But is there something I can replace the XXX with to tell it is a class object I am passing.

Answer 1

If you replaced XXX with SOMETHING_3.class_type.instance_type the code you've given should work.

It's worth noting that this becomes much more tricky (I believe it's currently impossible) to do if you're instead trying to take in an array of jitclass objects. If your full problem/code involves arrays of these jitclass objects, I would recommend you consider doing this with NumPy structured arrays rather than a jitclass. This is mainly because using an array of jitclass objects as a function parameter doesn't seem to to be supported in the current version of Numba. The reason for this is that an array of jitclass objects would be interpreted as a NumPy array with dtype of numpy.object, which is not a supported dtype in Numba's nopython mode. Since it's the type that Numba would be unable to lower (compile for use in nopython mode), nopython mode will fail for both lazy compilation (no function signature), and eager compilation (specifying function signature).

UPDATE:

Lists of jitclass objects are now supported, but there is very large overhead in passing them between Python and nopython compiled code (as of writing this), so keep that in mind.