如何用Numba实现(N,N)与(N,M,O)矩阵沿O维度高效相乘？

问题描述 #

需要用Numba即时编译函数，实现尺寸为$(N,N)$的矩阵Pi与尺寸为$(N,M,O)$的矩阵X沿O维度相乘（即左乘X沿O维度的所有“切片”）。尝试两种方案后遇到问题：

方案1及警告 #

代码：

@njit
def fast_expectation(Pi, X):
    
    res = np.empty_like(X)
    
    for i in range(Pi.shape[0]):
        for j in range(X.shape[1]):
            for k in range(X.shape[2]):
                res[i,j,k] = np.dot(Pi[i,:], X[:,j,k])
                            
    return res 

触发警告：

NumbaPerformanceWarning: np.dot() is faster on contiguous arrays, called on (array(float64, 1d, C), array(float64, 1d, A))

交换X维度后问题未解决。

方案2及错误 #

代码：

@njit
def multiply_ith_dimension(Pi, i, X):
    """If Pi is a matrix, multiply Pi times the ith dimension of X and return"""
    X = np.swapaxes(X, 0, i)
    shape = X.shape
    X = X.reshape(shape[0], -1)

    # iterate forward using Pi
    X = Pi @ X

    # reverse steps
    X = X.reshape(Pi.shape[0], *shape[1:])
    return np.swapaxes(X, 0, i)

触发错误：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
reshape() supports contiguous array only
...
    <source elided>
    shape = X.shape
    X = X.reshape(shape[0], -1)
    ^

解决方案 #

方法1：优化内存连续性，手动实现点积 #

警告核心是X[:,j,k]为非连续数组，Numba的np.dot对连续数组性能更优。手动展开点积运算，控制内存访问顺序：

@njit
def fast_expectation_optimized(Pi, X):
    N = Pi.shape[0]
    M = X.shape[1]
    O = X.shape[2]
    res = np.empty((N, M, O), dtype=X.dtype)
    
    for k in range(O):
        for j in range(M):
            for i in range(N):
                dot_sum = 0.0
                for n in range(N):
                    dot_sum += Pi[i, n] * X[n, j, k]
                res[i, j, k] = dot_sum
    return res

该方式避免调用非连续数组的np.dot，消除性能警告，Numba可充分优化嵌套循环。

方法2：调整数组内存布局，修复reshape错误 #

方案2的问题是swapaxes后的数组非连续，Numba的reshape仅支持连续数组。先转成连续数组再操作：

@njit
def multiply_ith_dimension_fixed(Pi, i, X):
    X_swapped = np.swapaxes(X, 0, i)
    # 转为连续数组，确保reshape可正常执行
    X_contiguous = np.ascontiguousarray(X_swapped)
    shape = X_contiguous.shape
    
    X_reshaped = X_contiguous.reshape(shape[0], -1)
    X_result = Pi @ X_reshaped
    
    result_reshaped = X_result.reshape(Pi.shape[0], *shape[1:])
    return np.swapaxes(result_reshaped, 0, i)

np.ascontiguousarray会在必要时复制数组，保证内存连续，解决reshape的类型错误。

方法3：利用Numba广播式矩阵乘法 #

若Numba版本较新（支持广播矩阵乘法），可直接简化操作，同时确保输入数组连续：

@njit
def fast_expectation_broadcast(Pi, X):
    # 确保X为连续数组，避免性能损耗
    X_contiguous = np.ascontiguousarray(X)
    # Pi(N,N)与X_contiguous(N,M,O)自动广播到每个O维度切片
    return Pi @ X_contiguous

该方式代码最简洁，Numba对矩阵乘法的优化成熟，性能最优。

性能建议 #

• 优先使用方法3，代码简洁且性能最佳；
• 若使用旧版Numba，选择方法1，手动循环经Numba优化后性能接近矩阵乘法；
• 所有场景下尽量保证输入数组为C连续（numpy默认布局），避免内存不连续导致的性能损失或错误。

内容的提问来源于stack exchange，提问作者Mr. Fafa