python - 如何在类的成员函数上使用 numba？

Question

我使用的是稳定版 Numba 0.30.1。

我能做到:

import numba as nb
@nb.jit("void(f8[:])",nopython=True)                             
def complicated(x):                                  
    for a in x:
        b = a**2.+a**3.

作为测试用例，加速是巨大的。但是如果我需要加速类中的函数，我不知道如何进行。

import numba as nb
def myClass(object):
    def __init__(self):
        self.k = 1
    #@nb.jit(???,nopython=True)                             
    def complicated(self,x):                                  
        for a in x:
            b = a**2.+a**3.+self.k

我为 self 对象使用什么 numba 类型？我需要将这个函数放在一个类中，因为它需要访问一个成员变量。

Answer 1

我遇到了非常相似的情况，我找到了一种在类内部使用 Numba-JITed 函数的方法。

诀窍是使用静态方法，因为这种方法不会在将对象实例添加到参数列表之前调用。无法访问 self 的缺点是您不能使用在方法之外定义的变量。因此，您必须将它们从有权访问 self 的调用方法传递给静态方法。就我而言，我不需要定义包装器方法。我只需要将我想要进行 JIT 编译的方法拆分为两个方法。

对于您的示例，解决方案是:

from numba import jit

class MyClass:
    def __init__(self):
        self.k = 1

    def calculation(self):
        k = self.k
        return self.complicated([1,2,3],k)

    @staticmethod
    @jit(nopython=True)                             
    def complicated(x,k):                                  
        for a in x:
            b = a**2 .+ a**3 .+ k