numba 中的协程

我正在开发一些需要快速协程的东西，我相信 numba 可以加快我的代码速度。

这是一个愚蠢的例子：一个函数对其输入进行平方，并添加其被调用的次数。

def make_square_plus_count():
    i = 0
    def square_plus_count(x):
        nonlocal i
        i += 1
        return x**2 + i
    return square_plus_count

你甚至不能nopython=FalseJIT 这个，大概是由于nonlocal关键词。

但你不需要nonlocal如果你使用一个类来代替：

def make_square_plus_count():
    @numba.jitclass({'i': numba.uint64})
    class State:
        def __init__(self):
            self.i = 0

    state = State()

    @numba.jit()
    def square_plus_count(x):
        state.i += 1
        return x**2 + state.i
    return square_plus_count

这至少有效，但如果你这样做就会崩溃nopython=True.

有没有一个可以编译的解决方案nopython=True?

如果你无论如何都要使用状态类，你也可以使用方法而不是闭包（应该是非Python编译的）：

import numba

@numba.jitclass({'i': numba.uint64})
class State(object):
    def __init__(self):
        self.i = 0

    def square_plus_count(self, x):
        self.i += 1
        return x**2 + self.i

square_with_call_count = State().square_plus_count  # using the method
print([square_with_call_count(i) for i in range(10)])
# [1, 3, 7, 13, 21, 31, 43, 57, 73, 91]

然而，时间显示这实际上比纯 python 闭包实现要慢。我希望只要你不使用nonlocalnumpy-arrays 或在方法（或闭包）中对数组进行操作，这会降低效率！