PyTorch如何无循环实现批量张量类外积的外和运算

云上老周

2026-8-26

解决方案

PyTorch 原生支持该运算，无需循环，也完全不需要使用指数转外积、手动构造笛卡尔积这类存在数值缺陷或效率低下的方案，直接利用张量广播机制一行代码即可完成，性能和数值稳定性都是最优的。

实现逻辑

你需要的运算本质是批量版本的「外和」，和外积的维度对齐逻辑完全一致：

原始输入x、y形状均为(num_batches, d)
给x在最后一维追加长度为1的维度，调整为形状(num_batches, d, 1)
给y在第二维（序列维度的位置）追加长度为1的维度，调整为形状(num_batches, 1, d)
两个调整维度后的张量直接相加，PyTorch 会自动通过广播机制把长度为1的维度做逻辑扩展，最终输出形状为(num_batches, d, d)的结果，严格满足osum[b, i, j] == x[b, i] + y[b, j]的要求。

代码示例

import torch

# 构造测试输入
num_batches, d = 3, 5
x = torch.randn(num_batches, d)
y = torch.randn(num_batches, d)

# 方式1：用unsqueeze增维，可读性更好
osum = x.unsqueeze(-1) + y.unsqueeze(1)

# 方式2：用None索引增维，写法更简洁，运行效果和上面完全一致
# osum = x[..., None] + y[:, None, :]

# 验证结果正确性
print(torch.allclose(osum[2, 1, 3], x[2, 1] + y[2, 3]))  # 输出True

方案优势

无任何Python层面的循环，所有计算走PyTorch底层优化的算子，执行效率最高
直接做原生加法运算，不存在数值精度损失或者稳定性问题
广播机制不会提前复制张量产生冗余内存占用，内存效率远高于手动构造笛卡尔积的实现
天然支持两个输入最后一维长度不同的场景：如果x形状为(num_batches, d1)、y形状为(num_batches, d2)，上述代码无需修改，直接输出形状为(num_batches, d1, d2)的正确结果。