Tensor 广播机制

广播（broadcasting）是 PyTorch 的 shape 对齐协议：不同 shape 的张量无需显式复制即可参与同一运算。规则本身三条就能讲完，但它是 silent bug 的最大来源之一——错误的对齐不报错，只是算出错误的 shape。本文给出规则、stride 视角、显式 API 与典型陷阱。

广播规则

两个张量运算时，shape 从右对齐逐维检查：

两维 size 相等 → 匹配
一方为 1 → 沿该维虚拟扩张至相等
维度数不同 → 小 shape 左侧补 1 再检查

x = torch.randn(8, 5, 1)
y = torch.randn(   1, 6)
(x + y).shape  # (8, 5, 6)：补 1 → (8,5,1) vs (1,1,6)，逐维取 max

心算流程固定：右对齐 → 补 1 → 逐维取 max。任何一维既不相等也不为 1，立即报错 The size of tensor a (m) must match the size of tensor b (n) at non-singleton dimension k——报错里的维度号是右数的索引，读报错时从最右维数起。

虚拟扩张：stride = 0

广播不复制数据。被扩张的维 stride 置 0，寻址时无论该维下标走多远都停在原地：

a = torch.randn(1, 3)
b = a.expand(4, 3)   # shape (4,3)，stride (0,1)，共享同一 Storage
a[0, 0] = 99
b[3, 0]              # 99，stride 为 0 的维上"多行"实为一份数据

这与视图、转置同属零拷贝元数据操作，物理基础见 Tensor 与 Storage。推论：广播读数据免费，但结果张量是实打实分配的——(1000, 1000) + (1000, 1) 的加法输出 1000×1000 显存，与先 repeat 再加完全一样，广播省的只是输入的复制。

显式 API

操作	行为	内存
`unsqueeze` / `[None]`	插入 size=1 的维	零拷贝
`expand` / `broadcast_to`	把 size=1 的维虚拟扩到指定大小	零拷贝，只读
`broadcast_tensors`	多个张量按规则同时广播到公共 shape	零拷贝
`repeat`	沿各维真实铺贴复制	完整分配

v = torch.randn(1, 4)
v.expand(3, 4)    # ✅ stride (0,1)
v.expand(3, 5)    # ❌ 只能扩 size=1 的维
v.repeat(3, 1)    # 真实 (3,4)，内存 ×3
 
x = torch.randn(3)
x[:, None]        # (3,1) 列向量，等价 x.unsqueeze(-1)
x[None, :]        # (1,3) 行向量

expand 后接 contiguous()/clone()/参与运算都会实体化。需要真实内存布局（如传给只接受连续张量的 kernel）时才用 repeat，其余场景一律 expand。

常见陷阱

(n,) 一维张量永远从最右维开始匹配。这是最高频的坑：

a = torch.randn(4, 3)
b = torch.randn(4)      # 意图：每行减一个标量
a - b                   # ❌ 维 1：3 vs 4
a - b[:, None]          # ✅ (4,1) → (4,3)

(n,)、(n,1)、(1,n) 三者语义完全不同：第一个是向量，后两个是矩阵。跨维广播前先 view(-1, 1) 或 [None] 明确落点。

归一化忘 keepdim。mean/sum/max 默认压缩维度，广播对齐就从错误的位置开始：

x = torch.randn(4, 3, 2)
x - x.mean(-1)                  # ❌ (4,3,2) vs (4,3)：2↔3 不匹配，报错
x - x.mean(-1, keepdim=True)    # ✅ (4,3,2) vs (4,3,1) → (4,3,2)

若尾部维数恰好相等（如方阵），不报错但算错——比报错更危险。

in-place 只能广播读取方。_ 后缀运算的写入方 shape 不能变：

x = torch.zeros(3, 1)
y = torch.randn(3, 4)
x.add_(y)   # ❌ 结果 (3,4) 装不进 x
y.add_(x)   # ✅ x 广播到 (3,4)

matmul 的批量广播。@ 把前导维当 batch，最后两维做矩阵乘，batch 维间正常广播：

h = torch.randn(4, 10, 768)   # (B, T, d)
w = torch.randn(768, 768)     # (d, d) 左侧补 1 → 广播
(h @ w).shape                 # (4, 10, 768)

实战模式

Attention 掩码。padding mask 是一维 (T,)，广播到任意 batch/head 数：

scores = q @ k.transpose(-2, -1) * d ** -0.5   # (B, H, T, T)
mask = valid[:, None, None]                     # (B,1,1) → key 维
scores = scores.masked_fill(~mask, float("-inf"))

Diffusion 加噪系数。每个样本有自己的 $\overset{α}{ˉ}_{t}$ ，一维 (B,) 系数对齐到 (B,C,H,W) 图像：

x_t = sqrt_ab[:, None, None, None] * x0 \
    + sqrt_1m_ab[:, None, None, None] * noise

这类 [:, None, None, None] 是 diffusion 训练代码的标准姿势，推导见 DDPM的数学原理推导。Classifier-Free Guidance 中条件/无条件批次拼接后共享同一组系数，同样靠这一步广播。

Linear 层权重共享。(d, d) 权重对 (B, T, d) 激活的投影无需复制权重，matmul 自动广播——参数量不随 batch 增长的原因正在于此。

心算 shape 是使用 PyTorch 的基本素养：读任何张量代码，先在脑内跑一遍「右对齐 → 补 1 → 逐维取 max」。维度变换的系统梳理见 Torch中调整Tensor维度，线性代数层面的语义见向量、矩阵、张量核心运算总结。

知识花园

探索