python - 当 python 类 jitclass 包含自身 jitclass 类时，如何使其兼容？

Question

我正在尝试创建一个类，它可以是 jitclass 的一部分，但具有一些本身是 jitclass 对象的属性。

例如，如果我有两个带有装饰器 @jitclass 的类，我希望在第三个类中实例化它们(组合)。

import numpy as np
from numba import jitclass
from numba import boolean, int32, float64,uint8

spec = [
    ('type' ,int32),
    ('val' ,float64[:]),
    ('result',float64)]

@jitclass(spec)
class First:
    def __init__(self):
        self.type = 1
        self.val = np.ones(100)
        self.result = 0.
    def sum(self):
        self.result = np.sum(self.val)

@jitclass(spec)
class Second:
    def __init__(self):
        self.type = 2
        self.val = np.ones(100)
        self.result = 0.
    def sum(self):
        self.result = np.sum(self.val)



@jitclass(spec)
class Combined:
    def __init__(self):
        self.List = []
        for i in range(10):
            self.List.append(First())
            self.List.append(Second())

    def sum(self):
        for i, c in enumerate(self.List):
            c.sum()
    def getresult(self):
        result = []
        for i, c in enumerate(self.List):
            result.append(c.result)
        return result


C = Combined()
C.sum()
result = C.getresult()
print(result)

在该示例中，我收到错误，因为 numba 无法确定 self.List 的类型，它是两个 jitclass 的组合。

如何使类 Combined 与 jitclass 兼容？

更新

它尝试了我在其他地方找到的东西:

import numpy as np
from numba import jitclass, deferred_type
from numba import boolean, int32, float64,uint8
from numba.typed import List

spec = [
    ('type' ,int32),
    ('val' ,float64[:]),
    ('result',float64)]

@jitclass(spec)
class First:
    def __init__(self):
        self.type = 1
        self.val = np.ones(100)
        self.result = 0.
    def sum(self):
        self.result = np.sum(self.val)
 

 
spec1 = [('ListA',  List(First.class_type.instance_type, reflected=True))]

@jitclass(spec1)
class Combined:
    def __init__(self):
        self.ListA = [First(),First()] 

    def sum(self):
        for i, c in enumerate(self.ListA):
            c.sum()
    def getresult(self):
        result = []
        for i, c in enumerate(self.ListA):
            result.append(c.result)
        return result


C = Combined()
C.sum()
result = C.getresult()
print(result)

但我收到此错误

List(First.class_type.instance_type)
TypeError: __init__() takes 1 positional argument but 2 were given

Answer 1

TL；博士:

• 您可以在 jitclass 中引用其他 jitclass，即使您有这些 jitclass 的列表。您只需要更正命名空间 numba.typed -> numba.types。
• 目前(从 numba 0.46 开始)不可能在 jitclasses 或非 Python numba.jit 函数中拥有异构列表。因此，您不能将 First 和 Second 的两个实例附加到同一个列表中。

解决numba.typed.List异常

您的更新几乎是正确的。您需要使用numba.types.List而不是numba.typed.List。区别有点微妙，但是 numba.types 包含签名类型，而 numba.typed 命名空间包含可以在代码中实例化和使用的类。

所以如果你使用它就会起作用:

spec1 = [('ListA',  nb.types.List(First.class_type.instance_type, reflected=True))]

更改此代码:

import numpy as np
import numba as nb

spec = [
    ('type', nb.int32),
    ('val', nb.float64[:]),
    ('result', nb.float64)
]

@nb.jitclass(spec)
class First:
    def __init__(self):
        self.type = 1
        self.val = np.ones(100)
        self.result = 0.
    def sum(self):
        self.result = np.sum(self.val)

spec1 = [('ListA',  nb.types.List(First.class_type.instance_type, reflected=True))]

@nb.jitclass(spec1)
class Combined:
    def __init__(self):
        self.ListA = [First(), First()] 
    def sum(self):
        for i, c in enumerate(self.ListA):
            c.sum()
    def getresult(self):
        result = []
        for i, c in enumerate(self.ListA):
            result.append(c.result)
        return result

C = Combined()
C.sum()
result = C.getresult()
print(result)

产生输出:[100.0, 100.0]。

Intermezzo:这里使用 jitclass 有意义吗？

但是这里要记住的是，普通的 Python 类可能会比 jitclass 方法更快(或同样快):

import numpy as np
import numba as nb

class First:
    def __init__(self):
        self.type = 1
        self.val = np.ones(100)
        self.result = 0.
    def sum(self):
        self.result = np.sum(self.val)

class Combined:
    def __init__(self):
        self.ListA = [First(), First()] 
    def sum(self):
        for i, c in enumerate(self.ListA):
            c.sum()
    def getresult(self):
        result = []
        for i, c in enumerate(self.ListA):
            result.append(c.result)
        return result

C = Combined()
C.sum()
C.getresult()

如果这只是出于好奇，那没问题。但对于生产，我会从纯 Python+NumPy 开始，仅在速度太慢时应用 numba，然后仅在成为瓶颈的部分上，并且仅当 numba 擅长优化这些东西时(numba 目前是专门的工具，而不是通用工具)。

使用 numba 的异构(混合类型)列表？

在 no-python(无对象)模式下使用 numba，您需要同类列表。据我所知，numba 0.46 不支持在 jitclasses 或 nopython-jit 方法中包含不同类型对象的列表。这意味着您不能有一个包含 First 和 Second 实例的列表。

所以这行不通:

self.List.append(First())
self.List.append(Second())

来自numba docs :

Creating and returning lists from JIT-compiled functions is supported, as well as all methods and operations. Lists must be strictly homogeneous: Numba will reject any list containing objects of different types, even if the types are compatible [...]