如何将numba与functools.reduce（）一起使用

Question

我有以下代码，我试图使用并行循环，并且：numbafunctools.reduce()mulimport numpy as npfrom itertools import productfrom functools import reducefrom operator import mulfrom numba import jit, prangelst = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]arr = np.array(lst)n = 3flat = np.ravel(arr).tolist()gen = np.array([list(a) for a in product(flat, repeat=n)])@jit(nopython=True, parallel=True)def mtp(gen):    results = np.empty(gen.shape[0])    for i in prange(gen.shape[0]):        results[i] = reduce(mul, gen[i], initializer=None)    return resultsmtp(gen)但这给了我一个错误：---------------------------------------------------------------------------TypingError                               Traceback (most recent call last)<ipython-input-503-cd6ef880fd4a> in <module>     10         results[i] = reduce(mul, gen[i], initializer=None)     11     return results---> 12 mtp(gen)~\Anaconda3\lib\site-packages\numba\dispatcher.py in _compile_for_args(self, *args, **kws)    399                 e.patch_message(msg)    400 --> 401             error_rewrite(e, 'typing')    402         except errors.UnsupportedError as e:    403             # Something unsupported is present in the user code, add help info~\Anaconda3\lib\site-packages\numba\dispatcher.py in error_rewrite(e, issue_type)    342                 raise e    343             else:--> 344                 reraise(type(e), e, None)    345     346         argtypes = []~\Anaconda3\lib\site-packages\numba\six.py in reraise(tp, value, tb)    666             value = tp()    667         if value.__traceback__ is not tb:--> 668             raise value.with_traceback(tb)    669         raise value    670 我不确定我哪里做错了。任何人都可以给我指出正确的方向吗？非常感谢。

Answer 1

您可以在 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