1 回答

TA貢獻1805條經驗 獲得超9個贊
您可以在 numba jitted 函數中使用 np.prod:
n = 3
lst = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
arr = np.array(lst)
flat = np.ravel(arr).tolist()
gen = [list(a) for a in product(flat, repeat=n)]
@jit(nopython=True, parallel=True)
def mtp(gen):
results = np.empty(len(gen))
for i in prange(len(gen)):
results[i] = np.prod(gen[i])
return results
或者,您可以使用如下所示的reduce(感謝@stuartarchibald指出這一點),盡管并行化在下面不起作用(至少從numba 0.48開始):
import numpy as np
from itertools import product
from functools import reduce
from operator import mul
from numba import njit, prange
lst = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
arr = np.array(lst)
n = 3
flat = np.ravel(arr).tolist()
gen = np.array([list(a) for a in product(flat, repeat=n)])
@njit
def mul_wrapper(x, y):
return mul(x, y)
@njit
def mtp(gen):
results = np.empty(gen.shape[0])
for i in prange(gen.shape[0]):
results[i] = reduce(mul_wrapper, gen[i], None)
return results
print(mtp(gen))
或者,因為Numba內部有一點魔力,可以發現將轉義函數并編譯它們的閉包。(再次感謝@stuartarchibald),你可以這樣,在下面:
@njit
def mtp(gen):
results = np.empty(gen.shape[0])
def op(x, y):
return mul(x, y)
for i in prange(gen.shape[0]):
results[i] = reduce(op, gen[i], None)
return results
但同樣,并行在numba 0.48之前在這里不起作用。
請注意,核心開發團隊成員推薦的方法是采用第一個使用 .它可以與并行標志一起使用,并具有更直接的實現。np.prod
添加回答
舉報