如何在Python标准库中实现批量维度特定的外积运算？

给定形状为(K, M)的矩阵X和(K, L)的矩阵V，需要对每个batch维度（第0维）的元素分别执行外积，最终得到形状为(K, M, L)的结果。直接使用np.outer会将数组展平，得到(K*M, K*L)的形状，无法满足需求。

高效实现方法 #

无需通过Python循环逐个处理batch，以下几种numpy原生方法可高效完成该操作：

1. 广播机制（最直观）

通过扩展维度让两个数组满足广播条件，直接相乘即可实现批量外积：

# 方法1：用np.newaxis快速扩展维度
W = X[:, :, np.newaxis] * V[:, np.newaxis, :]

# 方法2：用np.expand_dims更清晰指定扩展位置
W = np.expand_dims(X, axis=-1) * np.expand_dims(V, axis=1)

原理：将X扩展为(K, M, 1)，V扩展为(K, 1, L)，相乘时numpy会自动对每个batch的M维向量和L维向量做外积，最终输出(K, M, L)。

2. 爱因斯坦求和（语法简洁）

使用np.einsum通过维度标记直接定义计算逻辑：

W = np.einsum('km,kl->kml', X, V)

解释：km对应X的形状，kl对应V的形状，kml指定输出形状，表示对每个k（batch元素），将m和l维度的向量做外积。

3. 矩阵乘法（等价实现）

外积可以看作是列向量与行向量的矩阵乘法，通过扩展维度后使用np.matmul完成：

W = np.matmul(X[:, :, np.newaxis], V[:, np.newaxis, :])

这里X被转为(K, M, 1)的列向量，V转为(K, 1, L)的行向量，矩阵乘法后每个batch的结果就是对应向量的外积。

对比原循环实现 #

上述方法均为numpy向量化操作，相比原始的Python循环：

• 执行效率更高（底层由C实现，避免循环开销）
• 代码更简洁易读

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