torch.einsum 的基本语法是“下标字符串 + 张量列表”,下标中逗号分隔各输入张量维度标签,箭头后为输出标签;同一下标在多个输入中出现即求和,字母限a–z、大小写敏感,省略箭头时自动保留单次出现下标并按字母序排列。

PyTorch torch.einsum 的基本语法怎么写才不报错?
torch.einsum 的核心是“下标字符串 + 张量列表”,不是先写张量再凑下标。常见错误是维度顺序和下标字母数量对不上,比如把 torch.einsum('ij,jk->ik', A, B) 写成 torch.einsum('ij,jk->ik', A, B, C)(多传了张量)或 'i,j->ij'(逗号漏了,实际应为 'i,j->ij',但更常见的是 'i,j' 这种省略箭头的隐式模式——它会自动保留所有出现过的下标,且不求和)。
- 下标字符串中每个逗号分隔一个输入张量的维度标签,箭头后是输出维度标签
- 同一下标字母在输入中出现多次 → 表示沿该维度求和(如
'ij,jk'中j出现在两个张量里,自动求和) - 输出下标若省略(无箭头),则按隐式规则:保留所有只出现一次的下标,按字母序排列;若显式写出箭头,则必须严格匹配输出形状
- 字母只能是 a–z,不能用数字或下划线;大小写敏感,
'iI'是两个不同维度
哪些张量运算用 torch.einsum 替代原生函数更清晰?
矩阵乘、batched matmul、对角提取、转置、外积这些操作,用 torch.einsum 可读性反而更高,尤其当维度语义明确时(比如 b 代表 batch,h 代表 head,d 代表 dim)。
- 矩阵乘:
torch.einsum('ik,kj->ij', A, B)比A @ B更清楚谁是左乘谁是右乘,也避免torch.matmul对高维张量的广播逻辑干扰 - Batch 矩阵乘:
torch.einsum('bik,bkj->bij', Q, K)直观表达 batch 维度b不参与求和,比torch.bmm或torch.matmul少一层 reshape 需求 - 提取对角线:
torch.einsum('ii->i', X)比torch.diag(X)更通用(支持非二维张量,如'iijj->ij'提取块对角) - 外积:
torch.einsum('i,j->ij', a, b)比a.unsqueeze(1) * b.unsqueeze(0)更直白
torch.einsum 在性能和内存上有什么隐藏代价?
它不是万能加速器。PyTorch 会对简单下标(如 'ij,jk->ik')做优化路径,自动降级为 torch.matmul;但复杂下标(如含多个重复下标或非连续维度重排)可能触发通用路径,生成临时张量,导致显存暴涨或变慢。
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac,无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启:辅助功能 + 屏幕录制权限
- 避免在热循环里用动态构造的下标字符串(如 f-string 拼接),会阻止 JIT 编译和算子融合
- 如果只是转置或 reshape,优先用
.transpose()或.permute(),它们零拷贝;而torch.einsum('ijk->kij', x)会分配新内存 - 使用
torch.compile时,einsum支持有限,某些下标组合无法被图优化;实测中'bhqk,bhkd->bhdq'(Attention 中的 Scaled Dot-Product)能被良好优化,但带省略号(...)的下标目前仍不稳定
如何调试 torch.einsum 报错的维度不匹配?
最常遇到的是 RuntimeError: einsum(): operands do not match the given equation,本质是某个张量的实际维度数 ≠ 对应下标段的字母数,或某维度大小与同一下标在其他张量中不一致(除非该下标只在一个张量中出现,即不求和)。
立即学习“Python免费学习笔记(深入)”;
- 先打印每个张量的
shape和对应下标段长度,比如A.shape == (2, 3),下标'ij'就合法;若写成'ijk'就直接报错 - 若报错提示 “size mismatch for index ‘i’”,说明至少两个张量在
i维上 shape 不等,比如torch.einsum('ij,jk->ik', A, B)中A.shape[1] != B.shape[0] - 用
torch.einsum前可加断言:assert A.shape[1] == B.shape[0], f"dim mismatch: {A.shape[1]} vs {B.shape[0]}",比等报错再查更快 - 开发阶段建议开启
torch._dynamo.config.debug = True,它会在编译失败时输出更具体的下标推导过程
实际写的时候,别贪图一行写完所有逻辑;先拆成单步 einsum,验证中间形状,再合并。下标字符串看着像数学公式,但它本质是张量代数的 DSL——写错一个字母,结果就完全不对,而且往往没明显报错,只默默返回错误 shape。

















