Сопрограммы в Numba

Question

Я работаю над чем-то, что требует быстрых сопрограмм, и я считаю, 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=False JIT это, предположительно, из-за 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?

Answer 1

Если вы собираетесь использовать состояние класса в любом случае вы можете также использовать методы вместо закрытия (не должно быть никаких питон компилируется):

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. Я ожидаю, что до тех пор, пока вы не используете numpy-массивы nonlocal или не выполняете операции над массивами в своем методе (или закрытии), это будет менее эффективно!