如何嵌套 numba jitclass

Question

我试图了解 @jitclass 装饰器如何与嵌套类一起工作。我写了两个虚拟类：fifi 和 toto fifi 有一个 toto 属性。这两个类都有 @jitclass 装饰器，但编译失败。代码如下：

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

启动代码的脚本：

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])

当我运行脚本时，我得到 [...]

TypingError：未键入的全局名称“toto” 文件“fifi.py”，第 11 行

请注意，如果我删除 'fifi' 中对 'toto' 的任何引用，代码可以正常工作，并且由于 numba，我可以得到 x16 的加速。

Answer 1

可以将 jitclass 用作另一个 jitclass 的成员，尽管这样做的方式没有得到很好的记录。您需要使用 deferred_type 实例。这适用于 Numba 0.27 甚至更早版本。将fifi.py 更改为：

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

然后我得到输出：

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

这个功能可以在一些更高级的 jitclass 数据结构示例中看到，例如：

• stack.py
• linkedlist.py
• binarytree.py