PyTorch中(8,8)与(4,8,2)张量的广播矩阵乘法原理咨询

首先看你给出的张量运算示例：

a.shape 
=> (8, 8)
b.shape 
=> (4, 8, 2)

c = a @ b 
c.shape
=> (4, 8, 2)

你疑惑的核心是PyTorch中@运算符的批量矩阵乘法规则，它和普通张量的广播逻辑略有不同，具体计算过程拆解如下：

1. 维度对齐：自动扩展低维张量的前导维度 #

PyTorch的@运算符支持批量矩阵乘法，当两个张量维度数不同时，会自动给低维张量添加前导的批量维度，使其维度数与高维张量一致。这里：

• a是2维张量(8,8)，自动扩展为3维张量(1,8,8)（前导新增维度1）
• b是3维张量(4,8,2)

此时两个张量的维度为(1,8,8)和(4,8,2)，前导的批量维度1和4满足广播规则（1可以扩展为任意大小），因此可以进行批量运算。

2. 批量矩阵乘法的执行逻辑 #

扩展后，a的批量维度会被广播为4，相当于把原(8,8)的a复制4次，形成(4,8,8)的张量。然后对每个批量索引i（0到3）：

• 取广播后a的第i个张量：a_broadcasted[i] = a（形状(8,8)）
• 取b的第i个张量：b[i]（形状(8,2)）
• 执行标准矩阵乘法：a_broadcasted[i] @ b[i]，结果为(8,2)的张量

3. 结果堆叠 #

将4个批量运算得到的(8,2)张量在批量维度上堆叠，最终得到形状为(4,8,2)的张量c，和你看到的结果完全一致。

这个运算等价于手动循环计算每个批量的矩阵乘法再堆叠：

import torch
c = torch.stack([a @ b[i] for i in range(4)], dim=0)

关键区别：普通广播 vs 矩阵乘法广播 #

普通张量广播是针对元素级运算（如+、*），要求从最后一个维度往前匹配；而@的批量矩阵乘法是优先对齐维度数，扩展前导的批量维度，然后对每个批量执行标准矩阵乘法（要求前一个张量的最后一维等于后一个张量的倒数第二维，这里8=8满足条件）。

内容的提问来源于stack exchange，提问作者Andy Hin