消失的维度：keep_dims=false 为什么能把逆置换弄坏？

系列目录 · 数值与量化 · Read in English

“保持维度时全绿，去掉维度时全红。”这句话很像一个过于礼貌的报错：它几乎把答案告诉了你，却还是能让人对着屏幕坐很久。

求平均本身可能没有算错。真正出问题的是：归约前的第三根轴，在归约后已经不再有第三个位置，而恢复布局的代码还拿着旧地图。

一、轴不只是一个整数，它指向一份坐标约定

取教学张量 [A,B,C]=[4,6,9]，沿 B 归约。保留维度时，输出为 [4,1,9]；不保留时，输出为 [4,9]。

两种输出元素数量相同，但坐标系统不同。前者仍可用三个坐标描述，B 的坐标只能取 0；后者只有两个坐标。把二者混用，等于把一个城市的三维楼层地址塞进二维街道地图，再责怪导航软件没有找对门。

若硬件要求归约末轴，可以先将 [A,B,C] 变为 [A,C,B]。这时 B 从位置 1 移到位置 2。归约若保留轴，得到 [A,C,1]，逆置换后得到 [A,1,C]，最后移除单位维，才得到接口承诺的 [A,C]。

这条链的关键不是某个具体排列，而是中途保持同一秩：

1

逻辑输入 → 合法排列 → 保留归约维 → 逆排列 → 恢复接口形状

不能在“逆排列”之前擅自删除位置，因为逆排列仍然描述旧秩下的轴身份。

二、输出类型为什么不能直接拿来当中间类型

历史修复把归约过程中使用的类型，改为“按照有效输入形状将归约轴置一，再取原输出的元素类型”。这实际上把两个来源分开了：

• 中间形状来自当前坐标系及归约轴。
• 数值表示来自原算子的输出契约。

这两件事容易被一个“方便函数”一起复制。复制最终输出类型，看似省事，却会把最终的低秩形状过早带进中间节点。反过来，整个复制输入类型，又可能把输入的量化尺度带到归约结果。

因此，一个合适的类型构造接口应让调用者明确选择 shape 与 element type。若还有布局编码、存储空间等信息，也应说明是继承、重映射还是重新计算，不能笼统地称为“跟原来一样”。

用伪代码表达：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# order[new_axis] = old_axis
permuted_input = permute(input, order)
permuted_shape = [input_shape[old_axis] for old_axis in order]
mapped_axes = [new_axis for new_axis, old_axis in enumerate(order)
               if old_axis in normalized_axes]

kept_shape = copy(permuted_shape)
for axis in mapped_axes:
    kept_shape[axis] = 1

temporary_type = tensor(kept_shape, original_result_element_type)
reduced = native_reduce(permuted_input, axes=mapped_axes,
                        keep=true, result_type=temporary_type)
restored = permute(reduced, inverse(order))
result = reshape(restored, original_result_shape)

这个原则模型显式使用置换后的坐标系计算临时形状；若元素类型还携带逐轴量化轴等元数据，也需要相应重映射，示例未展开这部分。最容易犯的二次错误，是知道需要 keep=true，却依然用置换前的 axis 去改置换后的 shape。

三、四维变三维时，谁负责改轴？

假设另一个教学输入为 [1,4,6,9]，某阶段为了适配三维硬件，删除最前面的单位维。原 axis=2 对应的 B，在新表示里变成 axis=1。负轴也应先在原秩下归一化，再映射到新坐标系。

若图规范化阶段做过一次“减一”，算子专用合法化又按四维规则做一次“减一”，轴就走过站了。历史中确实出现过调整阶段顺序、把轴合法化移动到形状规范化之后的改动，其技术含义是建立单一责任：形状改变在哪发生，轴映射就应该在哪得到一致记录；后续阶段消费新的约定。

这里有三个不同的问题，不能用一个空列表含糊处理：

1. 原始语义就是对所有轴归约。
2. 原始只归约一个单位轴，这个轴后来被删除。
3. 原始有多个归约轴，其中一部分被删除，另一部分仍然存在。

第二种如果变成普通空 axes，而空 axes 被定义为全轴，就会从“值不变”变成“整个张量压成一个值”。历史中增加了记录被移除归约轴的标记，正是为了保存这个区别。恒等归约与重量化 将继续讨论它的量化含义。

四、排列正确，还得是硬件支持的排列

理论上的任意置换不等于后端已支持任意置换。历史差异中，一些循环排列被分解成可用的相邻轴交换，随后相关实现又在其他阶段演进。不能拿某一时期的测试形态，当成整个仓库永恒不变的接口。

这里的通用问题是：合法化不仅要把归约变成合法操作，也必须保证自己生成的辅助操作合法。如果“修复”一个不支持的归约时，顺手生成一个不支持的置换，错误只是晚了一站出现。

一种清晰的契约是：合法化后的每个节点都属于已声明的支持集合；每个新节点的类型能够由它的输入与属性独立推导；最终类型严格等于原算子输出类型。后端可以再做成本优化，但不能靠猜测补语义。

五、性能讨论要沿数据量走

置换前，张量有 A×B×C 个元素；归约后只有 A×C 个。因此，输入侧置换可能远比输出侧逆置换昂贵。简单地数“前后各一次置换”，会掩盖两者数据规模相差 B 倍的事实。

若每元素 b 字节，且两次置换都必须实际读写，可以用如下流量粗估：

1

额外流量 ≈ 2*b*A*B*C + 2*b*A*C

缓存、融合、布局视图、分块都会改变实际流量。它的作用是定位值得优先研究的部分：能否让生产者直接输出合法布局？能否让消费者接受逆置换前的布局？若不能，是否存在跨步归约替代方案？

即使消掉了一个 reshape，也未必节省数据移动，因为 reshape 可能本来只是视图。优化收益应该用数据搬运、指令和缓冲区生存期解释。

六、怎样写一个不会“太懂事”的测试

对称维度是这种错误的好朋友。若 A、B、C 都相等，很多错误排列仍有相同类型；若数据全部一样，错误归约轴也会给出相同答案。

建议用坐标编码输入，例如 x[a,b,c]=100*a+10*b+c，并让三个维度不同。这样错轴会留下清晰的数字指纹。系数应按维度范围设计，避免发生重叠；此处仅是小规模教学编码。

历史新增过针对不保留维度的 IR 测试，这能证明回归目标明确。读这样的补丁时，值得多问一句：测试检查的是最终 shape，还是还检查中间置换和归约的关系？前者防止接口错误，后者帮助定位错误发生在哪一步。

“消失的维度”并没有离开计算。它只是从数据结构里消失了。只要后面还有人在引用它，编译器就得给它留下一张可靠的转寄地址。

七、用一张轴身份表阻止第二次误修

调试时可以给轴取符号名，而不是一开始就记数字。设原轴依次叫 A、B、C。某个置换后的列表是 C、A、B；这个列表说明新位置零来自旧轴 C，新位置一来自旧轴 A。要沿旧轴 B归约，就在新列表里找 B，它现在位于二。逆置换则恢复 A、B、C的顺序，而不是把原排列数组再用一遍。

这区分了两种常见数组：一份表示“新位置取哪个旧位置”，另一份表示“旧位置搬到哪个新位置”。两者互为逆关系，只有某些交换排列碰巧相同。若测试只使用交换两个轴，代码即使混用两种约定也可能不失败。加入一次三轴循环排列，才会迫使二者显出区别。

把归约后的单位轴继续写成 B，而不是直接从列表删掉，可以看到整个过程中轴身份始终存在：C、A、B变成 C、A、B的一号长度版本，再逆排回 A、B、C，最后由外部接口决定删除 B。这就是保留维度策略为何容易证明，而不是某种形式上的偏好。

还要注意多个算子共享输入的情况。形状规范化改变一个值的类型时，所有使用者都在同一新坐标系中观察它。若只修了归约的 axes，却漏掉另一个轴敏感使用者，局部归约测试可能通过，完整图仍然错误。因此组合测试应把一个归约节点和另一个轴敏感节点同时接在同一输入上，验证公共映射只应用一次且都被消费。

这类错误适合用阶段快照定位：在规范化前后、合法化前后分别列出值的形状、轴身份和输出契约。不要把全部 IR混在一份巨大日志里凭直觉看。四张短表通常能直接指出是哪一个阶段第一次违反了不变量。

返回系列目录 · 上一篇 · 下一篇

• ← Previous Post
• Next Post →