如何用Numpy张量与向量相乘优化循环代码？

先明确原循环的计算逻辑：

• 对k的第二维（索引i）的每个切片k[:,i,:]（形状(5,4)），计算与b（形状(5,)）的点积，得到一个(4,)的数组
• 将该数组乘以dx后，加到y[i]上

你当前代码的问题在于错误构造B并误用矩阵乘法，导致维度不匹配。正确的向量化实现无需重复b，直接利用Numpy的广播和点积特性即可：

# 直接计算b与k的点积，得到形状为(3,4)的结果，对应每个i的delta_y
delta_y = b @ k
delta_y *= dx
y += delta_y
return y

原理说明 #

b是形状(5,)的一维数组，k是形状(5,3,4)的三维数组。使用@运算符时，Numpy会自动将b视为(1,5)的二维数组，与k的第一维（长度5）进行点积运算，最终得到形状为(3,4)的数组——这正好对应原循环中每个i生成的delta_y组成的集合。

如果需要更明确的维度控制，也可以用tensordot实现相同效果：

delta_y = np.tensordot(b, k, axes=1)  # 指定对b的第0维和k的第0维做点积
delta_y *= dx
y += delta_y
return y

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