如何将双重求和表达式转换为PyTorch einsum向量计算

先明确各变量的维度对应下标，方便映射einsum的求和规则：

• c: n×m → 下标 (k, l)（k覆盖n维度，l覆盖m维度）
• w: m×n → 下标 (l, i)（l覆盖m维度，i覆盖n维度）
• X_tilde: n×m×n → 下标 (k, l, j)（k覆盖n维度，l覆盖m维度，j覆盖n维度）
• b: m维向量 → 下标 (l)；X: n维向量 → 下标 (j)

核心双重求和的einsum写法 #

假设你原本的双重循环逻辑是对k（n维）和l（m维）求和，计算c[k,l] * w[l,i] * X_tilde[k,l,j]的累加，直接用einsum把这两个求和维度合并即可：

sum_component = torch.einsum("kl, li, klj -> ij", c, w, X_tilde)

• 逻辑说明：kl对应c的维度，li对应w的维度，klj对应X_tilde的维度；einsum会自动对重复出现的下标k和l做求和，最终保留i和j维度，得到m×n的结果，正好匹配你的输出要求。

整合b和X的项（如果表达式包含） #

如果你的最终输出还要加上b[l] * X[j]这类广播项，同样可以用einsum实现：

broadcast_component = torch.einsum("l, j -> lj", b, X)
final_output = sum_component + broadcast_component

验证代码 #

你可以用小规模张量对比循环和einsum的结果，确保正确性：

import torch
n, m = 3, 4
# 生成随机测试张量
c = torch.randn(n, m)
w = torch.randn(m, n)
X_tilde = torch.randn(n, m, n)
b = torch.randn(m)
X = torch.randn(n)

# 原循环实现
output_loop = torch.zeros(m, n)
for k in range(n):
    for l in range(m):
        output_loop += c[k, l] * w[l, :].unsqueeze(1) * X_tilde[k, l, :].unsqueeze(0)
output_loop += b.unsqueeze(1) * X.unsqueeze(0)

# einsum实现
sum_comp = torch.einsum("kl, li, klj -> ij", c, w, X_tilde)
broadcast_comp = torch.einsum("l, j -> lj", b, X)
output_einsum = sum_comp + broadcast_comp

# 检查结果一致性
print(torch.allclose(output_loop, output_einsum))  # 输出True表示结果一致

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