从 reshape 到 einops：理解 Transformer 代码中的 Tensor Shape DSL

起因

#

最近在做 CS336（Language Modeling from Scratch）的 Transformer 实现，代码里反复出现这种写法：

1rearrange(X, "... seq (heads d) -> ... heads seq d", heads=self.num_heads)

第一次看到时，它甚至比 x.view(...).transpose(...) 更加陌生——至少后者你知道它是 PyTorch 原生 API。

但看了一段时间以后会发现，它在解决 Transformer 代码中一个非常现实的问题：

Tensor 的 shape 越来越复杂，而数字 index 已经无法很好地表达每个 dimension 的语义。

Transformer 里经常出现的 shape：

1[B, S, D]
2[B, S, H, D_head]
3[B, H, S, D_head]
4[B, H, Q, K]

对于写代码的人来说，x.transpose(1, 2) 可能很清楚——因为他知道 dim 1 是 head，dim 2 是 sequence。

但对于读代码的人来说：

1dim 1 是 sequence？还是 head？
2dim 2 是 head？还是 hidden dimension？

必须根据上下文逐行推导。

这就是 einops 试图解决的问题。

einops 这个名字是什么意思

#

先把名字讲清楚，因为名字本身包含了核心思想。

根据 einops 官方 README：

einops stands for Einstein-Inspired Notation for operations (though “Einstein operations” is more attractive and easier to remember).

Notation was loosely inspired by Einstein summation (in particular by numpy.einsum operation).

关键词是 Einstein-Inspired Notation。

这个联系不难理解。看一个 attention score 的计算：

1scores = einsum(Q, K, "... query d_k, ... key d_k -> ... query key")

这里的 query、key、d_k 不是第 0、第 1、第 2 个 axis，而是具有语义的符号。d_k 同时出现在 Q 和 K 中表示要对这个 axis 做 reduction——这就是 Einstein summation 的核心思想：用 axis name 而不是位置 index 来描述 Tensor 操作。

einops 做的事情，就是把这种"通过名字描述 axis"的思想，从 einsum（多 Tensor 的 contraction）扩展到了单 Tensor 的 shape transformation：reshape、transpose、split、merge、reduction、repeat。

不需要从数学史去理解 Einstein summation convention。需要记住的只有一点：

einops 的 notation 来源于 einsum，核心想法是用 axis name 代替数字 index 来描述 Tensor 操作。

Mental model：把 pattern 看成 shape equation

#

这一节很重要。

不要把：

1rearrange(x, "b s (h d) -> b h s d", h=num_heads)

当作一种需要记住的新 API 调用方式。

更好的理解是：把 einops pattern 看成一个 Tensor shape equation。

1b s (h d)   →   b h s d

左边描述输入的 shape 语义：

1batch × sequence × (heads × head_dim)

右边描述期望的输出：

1batch × heads × sequence × head_dim

einops 负责根据这个声明完成实际的数据布局变化。

传统 PyTorch 写法：

1x = x.view(B, S, num_heads, D // num_heads)
2x = x.transpose(1, 2)

这描述的是"先 reshape，再 transpose"——一系列操作步骤。

einops 描述的是"输入是什么，输出是什么"——一个 shape transformation 的声明。

einops 更关注"输入是什么、输出是什么"，而不是"先 reshape 再 transpose"。

换一个角度：前者接近 imperative 风格（怎么做），后者接近 declarative 风格（要什么）。

一个容易产生的误解：axis name 不是 Tensor 的永久属性

#

这里有一个很容易产生的误解。

看到：

1rearrange(x, "batch sequence hidden -> batch hidden sequence")

很容易以为 x 从此拥有了 batch、sequence、hidden 这些命名的 dimension。

但实际上不是这样。这些名字只是 当前这一次 pattern 中的局部变量。

你完全可以写：

1rearrange(x, "foo bar baz -> foo baz bar")

只要 shape algebra 合法（三个 axis 做 permutation），einops 一样可以执行。

因此：einops 不是给 Tensor 的维度永久命名，而是在一次 operation 中给 logical axis 命名。

这也意味着 einops 没有能力检查你是否把一个 head axis 错写成了 sequence。它验证的是 shape algebra（维度乘积是否匹配），而不是 semantic correctness。

这一点对理解 einops 的能力边界很重要——它提供的是操作级别的 semantic naming，不是 Tensor 级别的 named dimension metadata（PyTorch 曾有过 named_tensor 的尝试，但几乎没人在用）。

从最简单的对比开始

#

假设：

1x.shape == (batch, sequence, hidden)

需要把 sequence 和 hidden 交换。

PyTorch：

1x = x.permute(0, 2, 1)

einops：

1x = rearrange(x, "batch sequence hidden -> batch hidden sequence")

看到 permute(0, 2, 1)，读者能得到的信息：第 1 和第 2 个 axis 被交换了。

看到 "batch sequence hidden -> batch hidden sequence"，读者直接知道：sequence 和 hidden 被交换了。

区别在于：

数字 index 表达的是 axis 的位置；axis name 表达的是 axis 的语义。

当 Tensor 只有 2-3 个 axis 时，这个区别不大。但 Transformer 中 Tensor 经常有 4 个甚至更多 axis（batch、head、sequence、head_dim），数字 index 带来的认知负担会迅速增加。

括号：split 和 merge dimension

#

rearrange 最值得讲清楚的是括号。

Transformer 里最经典的场景：Multi-Head Attention 中需要把一个 hidden dimension 拆成 num_heads × head_dim。

输入：

1[B, S, hidden]     其中 hidden = num_heads × head_dim
2                   例如 768 = 12 × 64

需要变成：

1[B, num_heads, S, head_dim]

传统 PyTorch：

1B, S, D = x.shape
2x = x.view(B, S, num_heads, D // num_heads)
3x = x.transpose(1, 2)

einops：

1x = rearrange(x, "b s (h d) -> b h s d", h=num_heads)

括号 (h d) 的含义：输入中这个 axis 的 size 等于 h × d，把它拆成两个 logical axis。

然后 -> b h s d 说明输出的 axis 排列方式。

flowchart TD
    A["[B, S, H*D]
输入：batch × seq × hidden"]
    B["[B, S, H, D]
split hidden → heads × head_dim"]
    C["[B, H, S, D]
rearrange axes"]
    A -->|"(h d) 拆分"| B
    B -->|"axis 重排"| C

一行 einops 完成了整个 transformation，而且 pattern 本身就是最好的注释。

反向操作同样直观。attention 计算完成后：

1[B, H, S, D]  →  [B, S, H*D]

1x = rearrange(x, "b h s d -> b s (h d)")

右边的括号 (h d) 表示把 h 和 d 两个 axis merge 成一个。

记忆方式（作为 intuition，不是严格语法定义）：

• 左边的括号 → 输入中一个 axis 是多个 logical axis 的乘积，需要 split
• 右边的括号 → 把多个 axis merge 成一个

CS336 中的真实代码

#

以下代码来自 CS336 Assignment 2 的 Transformer 实现（cs336_basics/model.py）。

Multi-Head Attention 的 split 和 merge

#

 1# 投影后，把 hidden 拆成 num_heads × d_k
 2Q, K, V = (
 3    rearrange(X, "... seq (heads d) -> ... heads seq d", heads=self.num_heads)
 4    for X in (Q, K, V)
 5)
 6
 7# attention 计算完成后，把 heads 和 d_v merge 回去
 8attn_output = rearrange(
 9    attn_output, "batch heads seq d_v -> batch seq (heads d_v)"
10).contiguous()

这也是我第一次看到这段代码时比较困惑的地方：... 是什么意思？

... 在 einops pattern 中表示"任意数量的 leading axes"，类似 NumPy 的 Ellipsis。它让 pattern 可以同时处理有或没有 batch dimension 的情况。

Scaled Dot-Product Attention

#

 1d_k = K.shape[-1]
 2attention_scores = einsum(
 3    Q, K, "... query d_k, ... key d_k -> ... query key"
 4) / math.sqrt(d_k)
 5
 6attention_weights = softmax(attention_scores, dim=-1)
 7
 8return einsum(
 9    attention_weights, V, "... query key, ... key d_v -> ... query d_v"
10)

第一个 einsum：Q 和 K 在 d_k 上做 contraction，得到 [..., query, key] 的 attention score。

第二个 einsum：attention weights 和 V 在 key 上做 contraction，得到 [..., query, d_v] 的输出。

对应数学：

如果用传统写法：

1scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

读者需要知道 K 的最后两个 axis 分别是什么，才能理解 transpose(-2, -1) 在做什么。而 einsum 的 pattern 直接告诉你：d_k 被消掉了，query 和 key 保留下来。

Linear 层

#

1def forward(self, x):
2    return einsum(x, self.weight, "... d_in, d_out d_in -> ... d_out")

连 nn.Linear 的 forward 都用 einsum 实现。pattern 直接说明：输入的最后一个 axis d_in 和 weight 的第二个 axis d_in 做 contraction，输出是 d_out。

RoPE 中的 pair split

#

1x1, x2 = rearrange(x, "... (half_d xy) -> xy ... half_d", xy=2).unbind(0)

RoPE 需要把 embedding dimension 拆成交替的 pair。(half_d xy) 表示最后一个 axis 可以看成 half_d × 2，拆开后 xy 移到最前面用 unbind(0) 分成两半。

如果用传统写法，大概需要：

1x1 = x[..., 0::2]
2x2 = x[..., 1::2]

或者：

1x = x.view(..., half_d, 2)
2x1, x2 = x[..., 0], x[..., 1]

哪种更清楚取决于上下文，但 einops 的写法至少把"这个 axis 的结构是 half_d × 2"这个信息写进了代码。

rearrange / reduce / repeat

#

einops 的三个核心操作：

前面已经详细讲了 rearrange。reduce 和 repeat 各看一个例子。

reduce

#

1from einops import reduce
2
3x = reduce(x, "batch sequence hidden -> batch hidden", "mean")

sequence 在右边消失了，因此沿 sequence 做 mean reduction。

对比 PyTorch：

1x = x.mean(dim=1)

dim=1 只说了位置。"batch sequence hidden -> batch hidden" 说了语义。

repeat

#

1from einops import repeat
2
3mask = repeat(mask, "query key -> batch head query key", batch=B, head=H)

输入只有 [query, key] 两个 axis，输出多了 batch 和 head——这两个新 axis 通过 repeat 创建。

直觉：如果右边出现了左边不存在的 axis name，就是在创建新维度。

从点积到 einsum

#

前面已经反复出现 einsum，也提到"output 中没有的 axis 会被 reduction"。这个描述没错，但对于已经忘记线性代数的工程师来说，中间缺了一层直觉。

我最开始看到：

1torch.einsum("ik,kj->ij", A, B)

第一反应是：这不就是矩阵乘法吗？

后来又发现：

1torch.einsum("d,d->", a, b)

可以表示点积。

再看到：

1torch.einsum("ij->", A)

甚至可以做求和。

于是我反而更困惑了：einsum 到底是什么？为什么它能表示这么多东西？

先看点积

#

从最简单的开始。两个长度为 3 的向量：

点积：

拆成两步：

1对应位置相乘：  [1×4, 2×5, 3×6] = [4, 10, 18]
2
3然后求和：      4 + 10 + 18 = 32

如果把这个共同维度叫 d，那么：

1einsum("d, d -> ", a, b)

表达的就是：

1result = 0
2for d in range(D):
3    result += a[d] * b[d]

这里 d 出现在输入中，但不出现在 -> 右边的输出中。

\(\Sigma\) 并不神秘，可以暂时直接理解成一个 for-loop + 累加。

axis 消失意味着什么

#

现有文章里已经说过"右边没有的 axis 会被 reduction"。补充一个更直观的理解：

1输入：d, d
2输出：（空）

d 在输出中消失了。原因不是它被"删除"了，而是 a[0]*b[0]、a[1]*b[1]、a[2]*b[2] 这些值通过 + 合成了一个 scalar。

axis 消失 = 这个 axis 上的多个值被 summation 成了一个值。

为什么矩阵乘法也能用 einsum

#

矩阵乘法看起来比点积复杂，但本质很简单。

\(C[0, 0]\) 是什么？

1A 的第 0 行 [1, 2]  dot  B 的第 0 列 [5, 7]
2
3= 1×5 + 2×7 = 19

\(C[0, 1]\) 呢？

1A 的第 0 行 [1, 2]  dot  B 的第 1 列 [6, 8]
2
3= 1×6 + 2×8 = 22

矩阵乘法，本质上就是左矩阵的每一行，与右矩阵的每一列分别做点积。

翻译成 for-loop：

1for i in range(I):
2    for j in range(J):
3        C[i, j] = 0
4        for k in range(K):
5            C[i, j] += A[i, k] * B[k, j]

i 选 A 的行，j 选 B 的列，k 是这对向量内部做点积的维度。

因此：

对应：

1torch.einsum("ik,kj->ij", A, B)

einsum 的阅读规则

#

到这里可以总结一个 practical mental model：

1ik, kj -> ij

怎么读？

1i → 出现在 output → 保留（外层 loop）
2j → 出现在 output → 保留（外层 loop）
3k → 只在 input 中出现 → reduction（内层 loop + 累加）

出现在 output 中的 index 是外层 loop（free index）；input 中出现但 output 中消失的 index 是 reduction loop（summation index）。

1# 对应的 loop 结构
2for i in ...:        # free
3    for j in ...:    # free
4        for k in ...:    # reduction
5            C[i,j] += A[i,k] * B[k,j]

这个规则适用于所有 einsum pattern。

三个 pattern 展示统一性

#

Batch matrix multiplication 没有引入新的数学概念——只是多了一个 free index b，对应多了一层外部 loop。

从 einsum 的视角看，这三者不是三套不同的 operation，只是 free index 和 reduction index 的组合不同。

Outer product：einsum 不一定 sum

#

为了避免"einsum 一定有 sum"的误解，看这个：

1einsum("i, j -> ij", a, b)

a = [1, 2]，b = [10, 20, 30]：

1[[10, 20, 30],
2 [20, 40, 60]]

i 和 j 都出现在 output 中，所以没有 axis 消失，没有 reduction。

只是产生所有 (i, j) 位置的乘积。

einsum 是否发生 summation，完全由某个 index 是否在 output 中消失决定。

Tensor contraction

#

到这里可以提一下术语了。上面这些操作的统一抽象叫 tensor contraction：

两个或多个 Tensor 按照共享的 logical index 对齐，在某些 index 上 multiply + summation，同时保留其他 index。

点积、矩阵乘法、batch 矩阵乘法、outer product——都是 contraction 的特例。

einsum 抽象的不是 matrix multiplication，而是更一般的 tensor contraction。

但 einsum 也不是万能的。它特别适合的是 index-based 的 multiply + sum 操作。不能自然表达 ReLU、softmax、sort、top-k、argmax 这类操作——这些需要的不是"沿 axis 乘加"，而是 elementwise nonlinearity 或 order statistics。

einsum 不等于"自动 transpose + matmul"

#

看这个例子：

1einsum("id, jd -> ij", A, B)

A shape 是 [I, D]，B shape 是 [J, D]。

如果用 @：

1A @ B    # 不行，D 和 J 不匹配
2A @ B.T  # 需要手动 transpose

但 einsum 不是"帮你自动 transpose 了"。更准确的理解是：

einsum 根本不是从"矩阵怎么摆才能 matmul"出发的。它直接描述 logical index 之间的关系。

它定义的是：

至于 backend 最终怎么执行（transpose、permute、GEMM kernel、fusion），那是 implementation detail。

这和文章前面讲 rearrange 的哲学一样：描述"算什么"，不是"怎么算"。

torch.einsum 和 einops.einsum

#

最后明确这两者的关系：

1# torch.einsum — 单字符 label
2torch.einsum("bhqd,bhkd->bhqk", Q, K)
3
4# einops.einsum — 完整 axis name，Tensor 在前
5from einops import einsum
6einsum(Q, K, "batch head query d, batch head key d -> batch head query key")

数学完全一样。einops.einsum 没有发明新的 einsum 数学。

它做的事情是把 einops 的 readable named-axis notation 带到了 einsum 上——和 rearrange 用完整名字代替位置 index 的思路一致。如果 operand 是 PyTorch Tensor，底层仍然用 PyTorch backend 执行。

重新看 CS336 的 attention

#

回到前面的代码：

1einsum(Q, K, "... query d_k, ... key d_k -> ... query key")

现在可以完全解构：

1query → output 中有 → free index（保留）
2key   → output 中有 → free index（保留）
3d_k   → output 中没有 → reduction index（消失 = summation）
4...   → 任意 leading axes → 全部保留

对应 loop 结构（忽略 ...）：

1for query in ...:
2    for key in ...:
3        scores[query, key] = 0
4        for d_k in ...:
5            scores[query, key] += Q[query, d_k] * K[key, d_k]

这时"在 d_k 上做 contraction"就不再是一句抽象描述，而是有了具体的 for-loop 直觉。

看到 einsum，不要先问"这是什么矩阵操作"，先问：哪些 index 保留，哪些 index 消失。

可读性真正解决的三个问题

#

到这里可以回答全文最核心的问题：einops 为什么在 Transformer 代码中更容易读懂？

数字 index 丢失语义

#

1x.permute(0, 2, 1, 3)

读者必须先建立 mental mapping：

10 = batch, 1 = sequence, 2 = head, 3 = head_dim

然后才能理解这行代码是在做什么。

1rearrange(x, "batch sequence head head_dim -> batch head sequence head_dim")

不需要这层 mental mapping。

reshape + transpose 描述的是操作步骤

#

1x = x.reshape(B, S, H, D).transpose(1, 2)

描述的是"先 reshape 再 transpose"——implementation procedure。

1x = rearrange(x, "b s (h d) -> b h s d", h=H)

描述的是"输入结构是什么，输出结构是什么"——shape intent。

前者更接近 implementation procedure，后者更接近 shape intent。

Pattern 是轻量的 shape contract

#

1rearrange(x, "b s (h d) -> b h s d", h=12)

如果 x 的最后一个 axis 不能被 12 整除，einops 会直接报错。

pattern 不仅仅是注释——它实际参与 runtime shape validation。这比代码注释强，因为注释不会随代码变化自动更新，而 pattern 会在 shape 不匹配时立即报错。

但也不要夸大：这不是一个完整的 type system，它只检查 shape algebra，不检查 semantic。

Trade-off：einops 不是银弹

#

对于：

1x.transpose(-1, -2)

这种非常简单、上下文明确的操作，直接用 PyTorch 完全没有问题。强行改成 einops 反而可能更啰嗦。

einops 的价值随着以下因素增加：

• Tensor rank 增加（4D、5D）
• axis 语义增加（batch、head、kv_head、sequence、head_dim、expert、token……）
• reshape / transpose / split / merge 的组合增加

尤其 Transformer、ViT、MoE、GQA/MQA 等结构中，大量 logical dimension 同时存在时，einops 的价值才特别明显。

einops 的价值不在于让所有 Tensor 操作都变短，而在于让复杂 Tensor transformation 保留足够多的语义信息。

一个简单的判断标准：如果你看到一行 permute(0, 2, 1, 3) 或 view(B, T, H, D).transpose(1, 2) 时，需要停下来想一想每个 index 对应什么——那这行代码可能值得用 einops 重写。

如果 transpose(-1, -2) 在上下文中已经足够清楚（比如 K.transpose(-1, -2) 在 attention 中几乎是 idiom），那就没必要改。

回到 CS336

#

学习 Transformer 的过程中，我慢慢发现一个规律：

很多时候真正让我看不懂一段代码的，并不是矩阵乘法本身，而是：

这个 Tensor 现在到底是什么 shape，每一个 dimension 又代表什么？

从 [B, S, D] 投影后变成 [B, S, H, D_head]，再 transpose 成 [B, H, S, D_head]，做完 attention 又变成 [B, H, Q, K]——如果代码里只有 reshape、transpose、permute，读者需要一直在脑中维护 axis index 和语义的映射。

而 einops 做的一件看起来很小的事情，就是：把这个映射直接写进代码。

因此我现在倾向于把 einops 理解成：

一种描述 Tensor shape transformation 的语言（DSL），而不仅仅是 reshape / transpose 的语法糖。

它的核心价值不在于某个 API 更短或更快，而在于：当 Tensor 的 shape 变得复杂时，它让代码的读者不再需要维护一张 “axis index → axis meaning” 的心理对照表。

参考资料

#

• einops GitHub — 官方 README 中的 Naming 和 “Why use einops notation” 部分是理解设计哲学的最佳入口
• Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation (ICLR 2022) — Alex Rogozhnikov 的论文，Section 4 详细定义了 notation 的形式语法
• Stanford CS336: Language Modeling from Scratch — Assignment 2 的 Transformer 实现中大量使用 einops
• PyTorch torch.einsum 文档 — 理解 Einstein summation 在 PyTorch 中的原生支持
• NumPy numpy.einsum 文档 — einops notation 最初的灵感来源