NumPy/PyTorch下批量张量按指定轴求和点乘的优雅实现

你之前调用np.tensordot得到错误形状的核心原因：仅指定了需要做求和的两个维度（A的axis=2、B的axis=1），没有约束两个张量共有的batch维度（长度为8的axis=0）需要逐位置匹配，导致tensordot将剩余维度做全笛卡尔积，最终输出多了交叉维度的错形状。

NumPy 实现 #

以下方法完全等价于手写循环的计算逻辑，均为经过BLAS优化的向量化实现，性能远高于Python层循环（尤其是batch维度较大时，循环的Python层开销会非常明显）：

• 批量矩阵乘法（内存开销最低）
  利用np.matmul的自动维度对齐规则：输入张量的最后两维会被识别为独立矩阵，前置维度按batch逻辑逐位置匹配。传入形状为(8,1,128)的A、(8,128,9)的B时，会自动完成8组独立的矩阵乘法，输出形状为(8,1,9)，压缩掉大小为1的冗余维度即可：

  C = np.squeeze(A @ B)
  

• 爱因斯坦求和（逻辑最直观）
  通过下标显式指定维度的匹配、求和规则，不会出现冗余维度错配的问题：

  # 下标规则：逐batch(b)匹配，对A的最后一维(k)和B的中间维(k)求和，输出剩余的batch和目标维度
  C = np.einsum('bik,bkj->bj', A, B)  # 直接输出形状(8,9)，无需额外维度压缩
  

PyTorch 实现 #

逻辑和NumPy完全对齐，GPU场景下的性能优势会更明显：

• 调用官方批量矩阵乘接口torch.bmm，天然支持3维张量的逐batch矩阵运算：

  C = torch.bmm(A, B).squeeze()
  

• 同样支持爱因斯坦求和写法，和NumPy版本的表达式完全通用：

  C = torch.einsum('bik,bkj->bj', A, B)
  

验证提示：可以用np.allclose(循环计算的C, 向量化计算的C)校验结果一致性，避免维度错配导致的计算错误。

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