python - 更新 numpy 矩阵的并行化函数

Question

假设我有一个函数，它接收一个 4D 矩阵和 4 个其他数组并更新矩阵条目：

@njit
def func(matrix, A, B, C, D):
    for i in range(len(A)):
        for j in range(len(B)):
            for k in range(len(C)):
                for l in range(len(D)):
                    matrix[i][j][k][l] = another_func(A[i], B[j], C[k], D[l])

的一个例子another_func是下面的形式：

def another_func(a, b, c, d)
    if c <= 5:
        return b
    else:
        return 1 + d * (c - a) + (b - d)

有没有办法并行化这些更新或加快速度？

我尝试使用 numba 的 @njit，但加速不够好。

标签： pythonpython-3.xnumpyparallel-processingnumba