Numpy中无循环实现3D数组与2D数组的对应维度矩阵乘法

当然有更简洁的实现方式！numpy.einsum和tensordot都是绝佳选择，完全不需要写for循环，而且代码可读性更强，能直接体现你的运算逻辑。

方法一：使用numpy.einsum（最直观） #

einsum的优势在于可以通过下标直接描述张量之间的运算关系，完全贴合你的需求：

import numpy as np

# 假设已经有A(s,m,m)和B(m,s)
C = np.einsum('ijk, ki -> ji', A, B)

下标解释： #

• ijk对应A的维度：i是第一个维度s，j和k是后面两个m维度（代表每个m×m矩阵的行和列）
• ki对应B的维度：k是第一个m维度，i是第二个s维度（代表每一列的向量）
• -> ji指定输出维度：j对应m（输出的行），i对应s（输出的列），正好匹配你要的(m,s)形状的C

这个表达式直接翻译了你要的逻辑：对每个i（第i个A矩阵和第i列B向量），计算A[i,:,:] @ B[:,i]，然后把所有结果按列拼接成C。

方法二：使用numpy.tensordot #

tensordot通过指定张量之间的收缩轴来实现批量运算，稍微需要调整一下维度顺序：

# 先做张量点积，再提取对角线维度
temp = np.tensordot(A, B, axes=([2], [0]))  # 形状为(s, m, s)
C = np.diagonal(temp, axis1=0, axis2=2).T  # 提取s维度的对角线，转置后得到(m,s)

逻辑解释： #

1. tensordot(A, B, axes=([2], [0]))：将A的第三个维度（m）和B的第一个维度（m）做收缩（点积），得到形状为(s, m, s)的中间张量，其中temp[i, j, k]代表A[i,j,:]和B[:,k]的点积结果。
2. np.diagonal(temp, axis1=0, axis2=2)：提取axis0（第一个s）和axis2（第二个s）的对角线元素，也就是取temp[i,j,i]，这正好对应A[i,j,:]和B[:,i]的点积，得到形状为(m, s)的结果。

对比你提到的重塑方法 #

你之前想到的重塑方法其实也可行，但代码稍微繁琐一点：

C = (A @ B.reshape(m, s, 1)).squeeze().T

这里把B重塑为(m,s,1)，让A的每个(m,m)矩阵和B的每个(m,1)向量相乘，得到(s,m,1)的结果，挤压掉最后一个维度后转置成(m,s)。这种方法也能工作，但einsum的代码更直接，一眼就能看懂运算逻辑。

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