Coroutines en numba

Estoy trabajando en algo que requiere corrutinas rápidas y creo que numba podría acelerar mi código.

Este es un ejemplo tonto: una función que ajusta su entrada y le agrega la cantidad de veces que se ha llamado.

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

Ni siquiera puede nopython=False JIT esto, probablemente debido a la palabra clave nonlocal .

Pero no necesitas nonlocal si usas una clase en su lugar:

 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 

Esto al menos funciona, pero se rompe si haces nopython=True .

¿Hay una solución para esto que compile con nopython=True ?

Si vas a usar una clase de estado de todos modos, también podrías usar métodos en lugar de un cierre (no debe comstackrse con 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] 

Sin embargo, los tiempos muestran que esto es en realidad más lento que una implementación de cierre de python puro. Espero que mientras nonlocal utilices matrices de direcciones no nonlocal realices operaciones con matrices en tu método (o cierre), ¡esto será menos eficiente!