当我使用 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

它说：

指定了关键字参数“parallel=True”，但无法进行并行执行转换。

该函数比上面演示的更复杂（numba 开销值得！）。另外，我**之前也尝试过电源，但结果是一样的。

我该如何修复它？

Answer 1

该消息仅说明您已指定，parallel但未发现任何可并行化的内容。这是因为np.power（这同样适用于数组幂运算符**）是在其他地方定义的函数 - 并且您无法并行化您正在调用的函数。但是，您可以并行化正在编写的函数或并行调用函数。

如果您想并行计算功率，您可以使用以下方法来实现numba.prange（另请参阅显式并行循环（numba 文档））：

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 并行化循环。