如何进行嵌套的 jitclass 数组实例化

Question

我正在尝试使用 numba 来加速我的代码，并且需要在另一个 jitclass 中使用一个 jitclass 的数组。然而，我似乎找不到将数组传递给 jitclass 的方法。

from numba import jitclass, int32

my_inner_spec = [
    ('a', int32)
]

@jitclass(my_inner_spec)
class MyInner:
  def __init__(self, a):
    self.a = a


my_outer_spec = [('inner', MyInner.class_type.instance_type[:])]

@jitclass(my_outer_spec)
class MyOuter:
  def __init__(self, inner):
    self.inner= inner

MyOuter([MyInner(0)])

这会产生此错误：

TypingError                               Traceback (most recent call last)
<ipython-input-106-3f644c4985b0> in <module>()
     18     self.inner= inner
     19 
---> 20 MyOuter([MyInner(0)])

...
...

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Internal error at <numba.typeinfer.CallConstraint object at 0x7fbdde552550>.
Failed in nopython mode pipeline (step: nopython mode backend)
Cannot cast reflected list(instance.jitclass.MyInner#7fbdde6c8650<a:int32>) to array(instance.jitclass.MyInner#7fbdde6c8650<a:int32>, 1d, A): %".19" = load {i8*, i8*}, {i8*, i8*}* %"inner"

File "<ipython-input-106-3f644c4985b0>", line 18:
  def __init__(self, inner):
    self.inner= inner
    ^

[1] During: lowering "(self).inner = inner" at <ipython-input-106-3f644c4985b0> (18)
[2] During: resolving callee type: jitclass.MyOuter#7fbdde6c8810<inner:array(instance.jitclass.MyInner#7fbdde6c8650<a:int32>, 1d, A)>
[3] During: typing of call at <string> (3)

Enable logging at debug level for details.

File "<string>", line 3:
<source missing, REPL/exec in use?>

Answer 1

正如 docs 所说，typed.List 是一项实验性功能，从普通 Python 列表推断类型似乎不起作用。

以下有点冗长，但它有效：

my_outer_spec = [('inner', nb.types.ListType(MyInner.class_type.instance_type))]

@nb.experimental.jitclass(my_outer_spec)
class MyOuter:
    def __init__(self, inner):
        self.inner = inner

lst = nb.typed.List()
lst.append(MyInner(0))
MyOuter(lst)

如何进行嵌套的 jitclass 数组实例化

How to do nested jitclass array instantiation

python

jit

numba