python - 当我使用 np.power() 时，如何在 numba 中设置 `parallel=True`？

Question

我将使用 numba 增强我的代码。然而，并行模式下的幂函数不能很好地工作，即对于以下函数:

import numpy as np
import numba

@numba.njit(parallel=True, fastmath = True)
def decay_rate(mV, mp):
    drate=(np.power(mp,-3))
    return drate

它说:

The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible.

该函数比上面演示的更复杂(numba 开销值得!)。另外，我之前尝试过使用 ** 来获取电源，但结果是相同的。

如何修复它？

Answer 1

该消息只是指出您指定了并行，但没有发现任何并行化的内容。这是因为 np.power (这同样适用于数组幂运算符 **)是在其他地方定义的函数 - 并且您无法并行化您正在调用的函数。但是，您可以并行化正在编写的函数或并行调用函数。

如果您想并行计算功率，可以使用 numba.prange (另请参阅 Explicit Parallel Loops (numba documentation) ):

import numba as nb
import numpy as np

@nb.njit(parallel=True, fastmath=True)
def decay_rate(mp):
    drate = np.empty(mp.size)
    for idx in nb.prange(mp.size):
        drate[idx] = np.power(mp[idx], -3)
    return drate

prange 明确告诉 numba 并行化循环。