如何用Numpy向量化实现Numpy矩阵的线性组合？

问题分析 #

直接使用*做元素乘法会触发广播错误，因为A形状为(3,5)，X形状为(5,4,4)，两者维度无法对齐广播。你需要实现的是用A的每行作为系数，对X中的5个4x4矩阵做线性组合，以下是几种高效的向量化实现方式：

方法1：使用tensordot（最直观） #

tensordot专门处理多维数组的缩并操作，指定axes=1让A的最后一维（长度5）和X的第一维（长度5）完成点积求和：

import numpy as np

np.random.seed(10)
X = np.random.rand(5,4,4)
A = np.random.rand(3,5)

result = np.tensordot(A, X, axes=1)
print(result.shape)  # 输出 (3,4,4)

方法2：广播+求和 #

给A增加两个长度为1的维度，让它与X的维度对齐后做元素乘法，再沿系数维度（axis=1）求和：

# 给A扩展维度为(3,5,1,1)，实现与(5,4,4)的X广播兼容
result = (A[:, :, np.newaxis, np.newaxis] * X).sum(axis=1)
print(result.shape)  # 输出 (3,4,4)

方法3：reshape后用矩阵乘法 #

把每个4x4矩阵拉平为1维向量，通过普通矩阵乘法完成线性组合后，再恢复原矩阵形状：

# X reshape为(5, 16)，A@X得到(3,16)，再reshape回(3,4,4)
result = (A @ X.reshape(5, -1)).reshape(3,4,4)
print(result.shape)  # 输出 (3,4,4)

以上三种方法均为完全向量化实现，效率远高于循环，且计算结果一致。

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