不同形状张量的组合逻辑、运算结果及形状预判方法（PyTorch）

一、不同形状张量的组合逻辑：广播机制 #

PyTorch中不同形状的张量能进行元素级运算（包括按位与），核心遵循广播机制，规则如下：

• 从张量的最后一个维度开始，向前逐一比对两个张量的维度大小
• 若某一维度的大小相等，或其中一个张量在该维度的大小为1，则视为匹配
• 若两个张量的维度数量不同，会自动给维度少的张量前面补1，直到两者维度数一致
• 满足匹配条件后，维度大小为1的张量会被逻辑上“扩展”（实际不复制内存，仅运算时复用）到与另一张量对应维度相同的大小，完成元素级运算

二、实例：张量A与B的按位与运算结果 #

给定：

• 张量A形状：(64, 1, 1, 42)
• 张量B形状：(1, 42, 42)

步骤1：对齐维度 #

B的维度数是3，比A少1，先给B前面补1，得到等效形状(1, 1, 42, 42)

步骤2：逐维度匹配与扩展 #

对比两个张量的每个维度（从前往后）：

1. 第1维（索引0）：64 vs 1 → B的该维度扩展为64
2. 第2维（索引1）：1 vs 1 → 无需扩展
3. 第3维（索引2）：1 vs 42 → A的该维度扩展为42
4. 第4维（索引3）：42 vs 42 → 完全匹配

最终结果 #

按位与运算后的张量形状为(64, 1, 42, 42)。运算过程中，A中每个(1,1,42)的子块会在第3维上扩展为(1,42,42)，再与B的对应元素做按位与；同时A的第1维有64个这样的子块，最终组合成64组(1,42,42)的结果。

三、提前判断运算后张量形状的方法 #

方法1：用PyTorch内置函数 #

直接调用torch.broadcast_shapes()，传入两个张量的形状即可得到广播后的结果形状：

import torch
shape_a = (64, 1, 1, 42)
shape_b = (1, 42, 42)
result_shape = torch.broadcast_shapes(shape_a, shape_b)
print(result_shape)  # 输出: torch.Size([64, 1, 42, 42])

方法2：手动推导 #

按照广播规则一步步来：

1. 补全维度：给维度少的张量前面补1，直到两者维度数相同
2. 逐维度取最大值（若其中一个维度大小为1，则取另一个的大小），最终得到的就是运算后的形状

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