问题描述

AddLayerNorm（及 InplaceAddLayerNorm）算子的 InferShape 推导 mean/rstd 输出形状时存在缺陷：仅对 x1 尾部第 1 根 norm 轴执行 SetDim(..., 1)，未覆盖 gamma 对应的全部尾轴。当 gamma 为多根尾轴时，mean/rstd 形状推导错误，并连锁导致 tiling 阶段 CheckX1GammaMean 校验失败，算子无法运行。

环境信息

• 仓库：cann/ops-nn
• 算子路径：norm/add_layer_norm/op_host/add_layer_norm_infershape.cpp
• 影响算子：AddLayerNorm、InplaceAddLayerNorm（共用同一 InferShape 函数）
• 触发条件：gamma 维数 > 1（即 norm 轴为多根尾轴，如 x1=[A,B,C,D]、gamma=[C,D]）

重现步骤

1. 构造输入 x1=[4,1,8,16]、gamma=[8,16]、beta=[8,16]。
2. 调用 AddLayerNorm 算子。
3. 观察 InferShape 输出的 mean/rstd 形状。

预期结果

mean/rstd 形状为 [4,1,1,1]（x1 尾部 gammaDimNum=2 维全部置 1）。

实际结果（缺陷）

mean/rstd 形状被错误推导为 [4,1,1,16]（仅置第 2 维为 1，漏置第 3 维），随后 tiling 侧 CheckX1GammaMean 因尾部维非 1 而报错失败。

日志/截图

问题代码（add_layer_norm_infershape.cpp:75）：

auto shape(*x1_shape);
shape.SetDim(shape.GetDimNum() - gamma_shape->GetDimNum(), 1);  // 仅置 1 根
*mean_shape = shape;
*rstd_shape = shape;

备注

• tiling 侧合轴逻辑、aveFactor、bias 分类、kernel 2D 扁平寻址均已正确支持多尾轴，仅 InferShape 为唯一缺陷点。
• 修复方案：改为循环置 1：

  for (size_t i = shape.GetDimNum() - gamma_shape->GetDimNum(); i < shape.GetDimNum(); i++) {
      shape.SetDim(i, 1);
  }
  

• gamma 为 1 维时循环只执行 1 次，行为与原实现完全一致，零回归。

关联 PR：https://gitcode.com/cann/ops-nn/pull/8804