Repository navigation
[Operator Mechanism] Reject non-empty index into a size-0 axis in gather/take_along_axis InferMeta - #79912
[Operator Mechanism] Reject non-empty index into a size-0 axis in gather/take_along_axis InferMeta#79912lxg5364 wants to merge 1 commit into
Conversation
Paddle-Bot Review Board （review完成)
可维护性建议（总分 90.5 / 100，只给建议，不影响评审结论）
各维度得分（仅评本 PR 改动，点击展开）
Powered by Nyanpasu claude with Opus 5.5 默认推理级别, please check the suggestions carefully. |
| int axis_v = axis < 0 ? axis + input_dim.size() : axis; | ||
| if (axis_v >= 0 && axis_v < input_dim.size() && input_dim[axis_v] == 0 && | ||
| IsKnownNonEmpty(index_dim)) { |
There was a problem hiding this comment.
test/legacy_test/test_take_along_axis_op.py 中 TestTakeAlongAxis0Size（第 35 行）和 TestTakeAlongAxis_ZeroSize（第 742 行）都使用 x=(2, 0, 5)、index=(2, 3, 5)、axis=1，并期望输出为全 0、可以求梯度。这正是这里新抛 OutOfRange 的输入，两个用例的 check_output/check_grad 都会报 IndexError。本 PR 只修改了 binary.cc，没有同步这两个用例。
建议：把这两个用例改为断言抛出 IndexError（覆盖动态图和 PIR），或删除它们，换成合法的 0-size 用例（例如 index 在 axis 上为 0）。
| if (!axis.FromTensor() && input_dim[axis_v] == 0 && | ||
| IsKnownNonEmpty(index_dims)) { |
There was a problem hiding this comment.
test_gather_op.py / test_take_along_axis_op.py 中补充 gather（动态图+静态图）和 take_along_axis 的单测，但本 PR 实际只改了 paddle/phi/infermeta/binary.cc，gather、take_along_axis、index_sample 三处检查都没有测试。
建议：在 test/legacy_test/test_gather_op.py、test_take_along_axis_op.py、test_index_sample_op.py 中补充用例，覆盖以下几类：
- 空轴 + 非空索引时抛出
IndexError； - 0-D index；
axis来自 Tensor 时跳过检查；- 空 index 仍正常返回。
同时请让 PR 描述与实际 diff 保持一致。
| PADDLE_THROW(common::errors::OutOfRange( | ||
| "The index is out of bounds: dimension [%d] of the input has size 0, " | ||
| "which leaves no valid index value, but received a non-empty index " | ||
| "of shape [%s].", | ||
| axis_v, | ||
| index_dims)); |
There was a problem hiding this comment.
💡 可维护性建议（AI 生成债务，不阻塞合入）：同一段“空轴 + 非空索引”判定和报错文本在 Gather（此处）、IndexSample（2638 行）、TakeAlongAxis（4627 行）复制了三份，并且已经出现不一致：IndexSample 用的是 errors::OutOfRange，另外两处用 common::errors::OutOfRange；axis 归一化也有三种写法，TakeAlongAxis 手写了 axis < 0 ? axis + size : axis，而本文件已引入 funcs::CanonicalAxis。
建议在 IsKnownNonEmpty 旁抽出一个 file-static helper，例如 CheckIndexOnEmptyAxis(input_dims, axis, index_dims)，在里面统一用 funcs::CanonicalAxis 归一化并抛 common::errors::OutOfRange，三个调用点各改成一行调用。
| static bool IsKnownNonEmpty(const DDim& dims) { | ||
| for (int i = 0; i < dims.size(); ++i) { | ||
| if (dims[i] <= 0) return false; | ||
| } | ||
| return true; |
There was a problem hiding this comment.
💡 可维护性建议（契约一致性，不阻塞合入）：注释说“只有每个 dim 都已知且为正时才返回 true”，但 DDim() 默认 rank_ = -1（paddle/common/ddim.cc:20），这时循环不执行，函数直接返回 true，未知 rank 会被当成“已知非空”。三个现有调用点都先约束了 rank，所以目前触发不到；以后复用这个 helper 时可能误报。
建议在循环前加 if (dims.size() < 0) return false;。0-D 标量应该仍返回 true，所以不要写成 <= 0。
…her/take_along_axis/index_sample InferMeta
d516bb6 to
cfcd104
Compare
| // Skipped when the axis comes from a Tensor, since its value is unknown at | ||
| // compile time. | ||
| if (!axis.FromTensor()) { | ||
| CheckIndexOnEmptyAxis(input_dim, axis_v, index_dims); |
There was a problem hiding this comment.
💡 可维护性建议（错误处理一致性，不阻塞合入）：这里用 axis.FromTensor() 跳过检查，但注释写的是“编译期未知”。动态图里 Tensor axis 的值已经能通过上方 axis.to<int>() 取到，却同样被跳过；而 kernel 侧也不会兜底：paddle/phi/kernels/cpu/gather_kernel.cc:50 在 x.numel() == 0 时直接返回，funcs/gather.h:173 的 GatherV2Function 也在逐元素越界校验之前提前返回。所以 paddle.gather(x, idx, axis=paddle.to_tensor(1)) 在空轴上仍会返回未初始化输出，和 int axis 抛 IndexError 不一致。
建议二选一：在 kernel 里 axis 归一化之后、提前返回之前补同样的 x.dims()[axis_v] == 0 && index.numel() > 0 检查；或者把注释和 CheckIndexOnEmptyAxis 的说明改成如实描述这条路径仍未覆盖。
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #79912 +/- ##
===========================================
Coverage ? 100.00%
===========================================
Files ? 1
Lines ? 13
Branches ? 0
===========================================
Hits ? 13
Misses ? 0
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
PR Category
Operator Mechanism
PR Types
Bug fixes
Description
When the gathered axis of
xhas size 0 andindexis non-empty, there is no validindex value: every index is out of bounds. PyTorch raises in this case, but Paddle did not:
paddle.gather, axis=0:GatherKernelallocatesoutand then returns early onx.numel() == 0, soout([N]) is returned uninitialized.paddle.gather, axis!=0:GatherV2CUDAFunctionreturns beforeResize/Alloc, so theforward "succeeds" with an unallocated output, and backward fails with
PreconditionNotMet: Tensor not initialized yet.paddle.take_along_axis:x.numel() == 0was silently zero-filled.This PR adds a host-side shape check in
GatherInferMetaandTakeAlongAxisInferMeta:when
x.dims()[axis] == 0andindexis known to be non-empty, raiseOutOfRange.Doing it in InferMeta keeps CPU/GPU/XPU and static graph consistent and avoids a
device-side assert. The check is skipped when
axiscomes from a Tensor or when theindex shape contains unknown (-1) dims.
Behavior change: these inputs used to return uninitialized values / zeros / an
unallocated tensor and now raise
IndexError. Valid empty cases are unchanged(empty index;
xempty on a non-gathered axis).Verification: unit tests for gather (dynamic + static) and take_along_axis are added
in
test_gather_op.py/test_take_along_axis_op.py. A 0-size sweep of 43 index-stylecases (getitem/setitem/gather/take_along_axis/put_along_axis) went from 18 paddle/torch
divergences to 0 for gather/take_along_axis.
paddle.index_samplehas the same problem: whenx.shape[1] == 0andindexisnon-empty there is no valid index value, but the op silently succeeded (PyTorch raises).
IndexSampleInferMetanow raisesOutOfRangein this case, using the same host-sidecheck as
gather/take_along_axis. Emptyindexand shapes with unknown (-1) dims arenot affected.
Verification:
index_sample(Tensor([3, 0]), Tensor([3, 2], int64))went from "Paddlesucceeds, Torch raises" to both sides raising
OutOfRange/out-of-bounds.是否引起精度变化
否
仅修改 InferMeta 的输入校验,未改动任何计算 kernel;合法输入的结果不变。此前"从空轴取
非空索引"这类本无合法结果的输入由返回未初始化值/0 改为显式报错。