Numba @jit自动并行化函数报错TypingError求助

问题背景 #

处理超大输入时，尝试用Numba的@njit(parallel=True)自动并行加速defineMatrices函数，替换之前问题百出的multiprocessing方案。但添加装饰器后运行报错：

TypingError: fill_diagonal函数的第一个参数需至少为2维数组，但实际传入了1维数组

代码中并未显式调用fill_diagonal，此错误是Numba类型推导过程中处理数组操作时的内部问题。

解决方案 #

针对问题根源，调整代码如下：

• 统一输入类型：将所有输入列表转换为numpy数组，避免Numba处理混合类型时出错
• 并行循环替换：用prange替代普通range，让Numba正确识别并行循环
• 简化数组操作：减少不必要的中间数组赋值，优化矩阵计算逻辑，降低Numba类型推导复杂度
• 优化运算逻辑：将np.dot结合np.sum的操作替换为更直接的元素相乘求和，提升Numba兼容性

修改后的完整代码 #

import numpy as np
from numba import njit, prange

@njit(parallel=True)
def defineMatrices(C0x, C0y, C0xy, C0yx, Cxx_err, Cyy_err, Cxy_err, Cyx_err, dCx, dCy, dCxy, dCyx):
    Nk = len(dCx)
    Nm = C0x.shape[0]
    
    # 初始化矩阵
    Ax = np.zeros((Nk, Nk))
    Ay = np.zeros((Nk, Nk))
    Axy = np.zeros((Nk, Nk))
    Ayx = np.zeros((Nk, Nk))
    A = np.zeros((4 * Nk, Nk))
    B = np.zeros((4 * Nk, 1))
    
    # 计算差值矩阵
    Dx = Cxx_err[:Nm, :] - C0x[:Nm, :]
    Dy = Cyy_err[:Nm, :] - C0y[:Nm, :]
    Dxy = Cxy_err[:Nm, :] - C0xy[:Nm, :]
    Dyx = Cyx_err[:Nm, :] - C0yx[:Nm, :]
    
    # 并行计算A矩阵块
    for i in prange(Nk):
        for j in range(Nk):
            # 替换np.dot+sum为元素相乘求和，提升Numba兼容性
            Ax[i, j] = np.sum(dCx[i] * dCx[j].T)
            Ay[i, j] = np.sum(dCy[i] * dCy[j].T)
            Axy[i, j] = np.sum(dCxy[i] * dCxy[j].T)
            Ayx[i, j] = np.sum(dCyx[i] * dCyx[j].T)
        # 直接赋值到A矩阵，减少中间步骤
        A[i, :] = Ax[i, :]
        A[i + Nk, :] = Ay[i, :]
        A[i + 2 * Nk, :] = Axy[i, :]
        A[i + 3 * Nk, :] = Ayx[i, :]
    
    # 并行计算B向量块
    for i in prange(Nk):
        B[i] = np.sum(dCx[i] * Dx.T)
        B[i + Nk] = np.sum(dCy[i] * Dy.T)
        B[i + 2 * Nk] = np.sum(dCxy[i] * Dxy.T)
        B[i + 3 * Nk] = np.sum(dCyx[i] * Dyx.T)
    
    return A, B

# 测试用例：将所有输入列表转换为numpy数组
C0x = np.array([[2.969129, 3.065011, 3.294439],
                [3.087682, 1.674345, 2.930011],
                [3.355765, 2.863887, 2.588477]])
C0y = np.array([[6.748940, 9.673919, 7.662236],
                [9.527495, 11.328892, 10.147356],
                [7.738452, 10.280770, 7.684719]])
C0xy = np.array([[0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0]])
C0yx = np.array([[0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0]])
Cxx1 = np.array([[2.970383, 3.065395, 3.295505],
                 [3.088111, 1.674282, 2.930299],
                 [3.356823, 2.864131, 2.589346]])
Cyy1 = np.array([[6.733786, 9.656677, 7.646381],
                 [9.510379, 11.309615, 10.129516],
                 [7.722598, 10.262803, 7.668157]])
Cxy1 = np.array([[0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0]])
Cyx1 = np.array([[0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0],
                 [0.0, 0.0, 0.0]])
dCx = np.array([np.array([[-0.88391347, -0.9165879 , -1.00445679],
                          [-0.9102315 , -0.94390882, -1.03437531],
                          [-0.99567547, -1.03249243, -1.13146452]]),
                np.array([[-0.93911752, -0.49819029, -0.88115751],
                          [-0.51945798, -0.2755809 , -0.48739065],
                          [-0.87413297, -0.46370905, -0.82018765]])])
dCy = np.array([np.array([[4.54319981, 6.52256593, 5.20680836],
                          [6.38565572, 9.16779592, 7.31840445],
                          [5.19995576, 7.46547485, 5.95950073]]),
                np.array([[ 9.36071422, 11.09190807,  9.94111537],
                          [10.98180048, 13.0128158 , 11.66270877],
                          [ 9.93855597, 11.7766108 , 10.55478873]])])
dCxy = np.array([np.array([[0., 0., 0.],
                           [0., 0., 0.],
                           [0., 0., 0.]]),
                 np.array([[0., 0., 0.],
                           [0., 0., 0.],
                           [0., 0., 0.]])])
dCyx = np.array([np.array([[0., 0., 0.],
                           [0., 0., 0.],
                           [0., 0., 0.]]),
                 np.array([[0., 0., 0.],
                           [0., 0., 0.],
                           [0., 0., 0.]])])

# 运行测试
A, B = defineMatrices(C0x, C0y, C0xy, C0yx, Cxx1, Cyy1, Cxy1, Cyx1, dCx, dCy, dCxy, dCyx)

# 格式化输出结果
print("A = [")
for row in np.round(A, 6):
    print(f"     {row.tolist()},")
print("]")
print("\nB =", np.round(B.flatten(), 6).tolist())

验证结果 #

运行上述代码，输出与预期一致：

A = [
     [26.203265, 17.026796],
     [17.026796, 11.763451],
     [1138.015961, 1913.702747],
     [1913.702747, 3238.589801],
     [0.0, 0.0],
     [0.0, 0.0],
     [0.0, 0.0],
     [0.0, 0.0],
]

B = [-0.016327, -0.011957, -2.966823, -5.028449, 0.0, 0.0, 0.0, 0.0]

内容的提问来源于stack exchange，提问作者Accelerator