PyTorch中如何对(n1,N)与(n2,N)的Tensor做广播得到(n1,n2)结果

你需要实现的是两组行向量的两两点积运算，最终输出的(n1,n2)张量中，第[i][j]位元素等于x1的第i行与x2的第j行逐元素相乘后求和的结果。

原代码错误原因 #

你原有代码中对x2的维度调整方向错误，reshape后的x2形状为(n2, N, 1)，和形状为(n1, N)的x1广播时维度无法对齐到两两行配对的逻辑，因此得到的结果不符合预期。

推荐实现方案 #

方案1：矩阵乘法实现（性能最优） #

矩阵乘法的运算逻辑完全匹配需求，仅需对第二个张量做转置后相乘即可：

import torch
# x1形状为(n1, N), x2形状为(n2, N)
result = torch.matmul(x1, x2.T)
# 也可以用@运算符简写
result = x1 @ x2.T

维度验证：(n1, N) @ (N, n2) = (n1, n2)，完全符合输出形状要求。

方案2：显式广播实现 #

如果需要明确用广播逻辑拆解运算过程，可以按如下方式调整维度：

# x1新增维度后形状为 (n1, 1, N)
x1_ext = x1.unsqueeze(1)
# x2新增维度后形状为 (1, n2, N)
x2_ext = x2.unsqueeze(0)
# 广播逐元素乘后，对最后一维求和
result = (x1_ext * x2_ext).sum(dim=-1)

该方案与矩阵乘法的计算结果完全等价，只是显式展示了广播配对的过程。

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