Is there a way to include a Jitclass in another Jitclass?

Numba Support: How do I do ...?
1

This question gets asked relatively often, here’s an example of one way to do it:

import numpy as np
from numba import int64
from numba.experimental import jitclass

spec_foo = {'data': int64[:]}

@jitclass(spec_foo)
class Foo(object):
    def __init__(self, init_data):
        self.data = init_data

# Want to have a Foo instance in the Bar class
spec_bar = {'foo_inst':  Foo.class_type.instance_type}

@jitclass(spec_bar)
class Bar(object):

    def __init__(self, foo_data):
        # This constructs the Foo instance in the Bar constructor
        self.foo_inst = Foo(foo_data)

# Want to have a Foo instance in the Baz clazz
spec_baz = {'foo_inst':  Foo.class_type.instance_type}

@jitclass(spec_baz)
class Baz(object):

    def __init__(self, foo_inst):
        # This assigns a reference to the Foo instance in the Bar constructor
        self.foo_inst = foo_inst

x = np.arange(5).astype(np.int64)
# Build a Bar instance
bar = Bar(x)
print(bar.foo_inst.data)

# Create a Baz instance that holds a reference to Bar's Foo instance.
baz = Baz(bar.foo_inst)

# Do some arithmetic...
baz.foo_inst.data += 1
print(baz.foo_inst.data)
print(bar.foo_inst.data)
1 Like
2

As of Numba 0.52 you can also use the new jitclass spec syntax.

import numpy as np
from numba import int64
from numba.experimental import jitclass

spec_foo = {'data': int64[:]}

@jitclass(spec_foo)
class Foo(object):
    def __init__(self, init_data):
        self.data = init_data
        
@jitclass
class Bar(object):
    foo_inst: Foo

    def __init__(self, foo_data):
        # This constructs the Foo instance in the Bar constructor
        self.foo_inst = Foo(foo_data)

@jitclass
class Baz(object):
    foo_inst: Foo

    def __init__(self, foo_inst):
        # This assigns a reference to the Foo instance in the Bar constructor
        self.foo_inst = foo_inst

x = np.arange(5).astype(np.int64)
# Build a Bar instance
bar = Bar(x)
print(bar.foo_inst.data)

# Create a Baz instance that holds a reference to Bar's Foo instance.
baz = Baz(bar.foo_inst)

# Do some arithmetic...
baz.foo_inst.data += 1
print(baz.foo_inst.data)
print(bar.foo_inst.data)
3

Thanks for sharing @EPronovost :slight_smile:

4

I need a jitclass that has takes as attribute a list of itself jitclasses. I’m trying this

@jitclass(spec=[("values", types.List(Bar.class_type.instance_type))])
class JitList(object):
    def __init__(self, values):
        self.values = values


FooList = JitList([bar])

but get the following error:

Failed in nopython mode pipeline (step: nopython mode backend)
Cannot cast ListType[instance.jitclass.Bar#2c5f1d61490<foo_inst:instance.jitclass.Foo#2c5efef29a0<data:array(int64, 1d, A)>>] to list(instance.jitclass.Bar#2c5f1d61490<foo_inst:instance.jitclass.Foo#2c5efef29a0<data:array(int64, 1d, A)>>)<iv=None>: %".68" = load {i8*, i8*}, {i8*, i8*}* %"$8call_method.3"

Any suggestions?

5

Hi @alanlujan91

I suggest writing a copy-pasteable reproducer so that someone can take a look at the code in more detail by running it locally. Perhaps even start a new topic with the reproducer?

Many thanks.

1 Like