How to nest numba jitclass

Question

I'm trying to understand how the @jitclass decorator works with nested classes. I have written two dummy classes: fifi and toto fifi has a toto attribute. Both classes have the @jitclass decorator but compilation fails. Here's the code:

fifi.py

from numba import jitclass, float64
from toto import toto

spec = [('a',float64),('b',float64),('c',toto)]

@jitclass(spec)
class fifi(object):
  def __init__(self, combis):
    self.a = combis
    self.b = 2
    self.c = toto(combis)

  def mySqrt(self,x):
    s = x
    for i in xrange(self.a):
      s = (s + x/s) / 2.0
    return s

toto.py:

from numba import jitclass,int32

spec = [('n',int32)]

@jitclass(spec)
class toto(object):
  def __init__(self,n):
    self.n = 42 + n

  def work(self,y):
    return y + self.n

The script that launches the code:

from datetime import datetime
from fifi import fifi
from numba import jit

@jit(nopython = True)
def run(n,results):
  for i in xrange(n):
    q = fifi(200)
    results[i+1] = q.mySqrt(i + 1)

if __name__ == '__main__':
  n = int(1e6)
  results = [0.0] * (n+1)
  starttime = datetime.now()
  run(n,results)
  endtime = datetime.now()

  print("Script running time: %s"%str(endtime-starttime))
  print("Sqrt of 144 is %f"%results[145])

When I run the script, I get [...]

TypingError: Untyped global name 'toto' File "fifi.py", line 11

Note that if I remove any reference to 'toto' in 'fifi', the code works fine and I get a x16 speed up thanks to numba.

JoshAdel · Accepted Answer · 2016-07-31T17:39:55.090

23

It is possible to use a jitclass as a member of another jitclass, although the way of doing this isn't well documented. You need to use a deferred_type instance. This works in Numba 0.27 and possibly earlier. Change fifi.py to:

from numba import jitclass, float64, deferred_type
from toto import toto

toto_type = deferred_type()
toto_type.define(toto.class_type.instance_type)

spec = [('a',float64),('b',float64),('c',toto_type)]

@jitclass(spec)
class fifi(object):
  def __init__(self, combis):
    self.a = combis
    self.b = 2
    self.c = toto(combis)

  def mySqrt(self,x):
    s = x
    for i in xrange(self.a):
      s = (s + x/s) / 2.0
    return s

I then get as output:

$ python test.py
Script running time: 0:00:01.991600
Sqrt of 144 is 12.041595

This functionality can be seen in some of the more advanced jitclass examples of data structures, for example:

edited Jul 31 '16 at 17:39

answered Jul 31 '16 at 14:18

JoshAdel

66,734
27
141
140

I was trying to use a jitclass as a member of another jitclass by using `typeof` on the instance after creating it, but this is much less awkward! – Chad Jan 24 '20 at 17:27
4

Note that the use of the deferred mechanism is only needed for self-referential classes, like the linked implementation in the examples. You can use `toto.class_type.instance_type` directly in the jitclass spec, which would then be: `spec = [('a',float64), ('b',float64), ('c',toto.class_type.instance_type)]`. From https://stackoverflow.com/questions/57640039/ – mgmalheiros Jan 04 '21 at 05:45
@mgmalheiros. Thank you. That is the correct answer – Mad Physicist Dec 06 '21 at 03:46

How to nest numba jitclass

1 Answers1

Linked