[0-size Tensor No.271] Add 0-size Tensor support for paddle.take_along_axis API.

[straigrand]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

本次提交为完成任务 #72637 中关于 paddle.take_along_axis 的部分。

修改历程介绍如下：

1. 问题复现与分析：

   • 根据任务要求，首先尝试使用PaddleAPITest复现BUG。在PaddleAPITest的--accuracy=True精度对比模式下，由于paddle.take_along_axis(arr, index, axis)与torch.take_along_dim(input, indices, dim)的参数名不统一，反复出现TypeError: missing a required argument的参数绑定错误，无法准确定位问题。
   • 切换到PaddleAPITest的--paddle_only=True模式后，成功触发了Python层的TypeError: take_along_axis() got an unexpected keyword argument 'index'报错，这暴露了API前后端参数名不一致的问题。
   • 为绕过Python接口问题，并直接验证C++ Kernel，最终采用unittest的OpTest框架编写单元测试。在未修复Kernel的情况下，成功复现了底层的C++错误：

     PreconditionNotMetError: Tensor holds no memory. Call Tensor::mutable_data firstly.
     [Hint: holder_ should not be null.] (at /home/aistudio/Paddle/paddle/phi/core/dense_tensor_impl.cc:46)
     [operator < take_along_axis > error]
     

   • 错误分析表明，当输入Tensor的元素个数为0时，算子Kernel未做特殊处理，导致程序在访问空Tensor的内存时崩溃。

2. 前向修复 (Forward Fix)：
   a. 定位API: 在Paddle/python/paddle/tensor/manipulation.py中找到了def take_along_axis(...)的Python定义，其核心实现调用了_C_ops.take_along_axis。
   b. 定位算子定义: 使用grep发现，该算子没有独立的.yml文件，其定义位于paddle/phi/ops/yaml/ops.yaml中。
   c. 检查InferMeta: 根据ops.yaml的指引，在paddle/phi/infermeta/binary.cc中找到了TakeAlongAxisInferMeta函数。经分析，其out->set_dims(index.dims())逻辑能正确推导0-size Tensor的输出形状，无需修改。
   d. 修改Kernel:
   * 根据grep结果，定位到CPU Kernel文件为paddle/phi/kernels/cpu/take_along_axis_kernel.cc。
   * 参照标准修复范式，在TakeAlongAxisKernel函数开头加入了对0-size情况的保护。核心逻辑是判断index.numel()是否为0，因为输出的形状完全由index决定。
   * 修复代码如下：

       if (index.numel() == 0) {
         dev_ctx.template Alloc<T>(out);
         return;
       }
       if (x.numel() == 0) {
         phi::Full<T, Context>(
             dev_ctx, common::vectorize(out->dims()), static_cast<T>(0), out);
         return;
       }

   d. 依照以上原则修改CPU、GPU、XPU Kernel

3. 反向修复 (Backward Fix)：

   • 本次修复的核心在于解决前向Kernel（take_along_axis_kernel）在处理0-size输入时，因访问未初始化内存而导致的PreconditionNotMetError。
   • 在添加的OpTest单元测试中，我们通过self.check_grad()对反向梯度进行了检查。
   • 在修复了前向Kernel之后，check_grad能够顺利通过。这表明，对于我们测试的0-size场景，正确的前向输出（一个0-size Tensor）能够保证反向计算图的正确建立和执行，因此本次并未对take_along_axis_grad_kernel.cc文件进行修改。

4. 添加单测 (Add Unit Test)：

• 在test/legacy_test/test_take_along_axis_op.py文件中，为彻底解决因父类TestTakeAlongAxisOp的setUp方法无法兼容0-size数据而导致的CI报错（IndexError, ValueError），最终方案是放弃继承，添加了两个全新的、独立的OpTest测试类，分别覆盖两种不同的0-size边界场景。

• 测试场景一：输入arr为0-size，但index不为0-size

  • 通过TestTakeAlongAxis0Size1类进行验证。完整测试代码如下：

    # --- 场景一: 输入x为0-size, index不为0-size ---
    class TestTakeAlongAxis0Size1(OpTest):
        # 修复 tearDownClass 的 op_type 缺失问题
        op_type = "take_along_axis"
    
        def setUp(self):
            self.python_api = paddle.take_along_axis
            self.op_type = "take_along_axis"
            self.dtype = "float64"
            self.check_pir = True
    
            x = np.zeros((2, 0, 5)).astype(self.dtype)
            indices = np.zeros((2, 3, 5)).astype("int64")
    
            self.inputs = {'Input': x, 'Index': indices}
            self.attrs = {'Axis': 1}
    
            # 手动定义期望的输出，不再依赖父类的复杂计算
            output = np.zeros((2, 3, 5)).astype(self.dtype)
            self.outputs = {'Result': output}
    
        def test_check_output(self):
            self.check_output(check_pir=self.check_pir)
    
        def test_check_grad(self):
            self.check_grad(['Input'], 'Result', check_pir=self.check_pir)

• 测试场景二：索引index为0-size，但arr不为0-size

  • 通过TestTakeAlongAxis0Size2类进行验证。完整测试代码如下：

    # --- 场景二: 索引index为0-size, x不为0-size ---
    class TestTakeAlongAxis0Size2(OpTest):
        # 修复 tearDownClass 的 op_type 缺失问题
        op_type = "take_along_axis"
    
        def setUp(self):
            self.python_api = paddle.take_along_axis
            self.op_type = "take_along_axis"
            self.dtype = "float64"
            self.check_pir = True
    
            x = np.random.rand(2, 3, 5).astype(self.dtype)
            indices = np.zeros((2, 0, 5)).astype("int64")
    
            self.inputs = {'Input': x, 'Index': indices}
            self.attrs = {'Axis': 1}
            
            output = np.zeros((2, 0, 5)).astype(self.dtype)
            self.outputs = {'Result': output}
    
        def test_check_output(self):
            self.check_output(check_pir=self.check_pir)
    
        def test_check_grad(self):
            # 输出为0-size，其梯度也为0-size。
            # 通过 user_defined_grads 手动指定反向输入的梯度，
            # 帮助框架完成梯度检查。
            grad = np.zeros_like(self.outputs['Result']).astype(self.dtype)
            self.check_grad(
                ['Input'],
                'Result',
                user_defined_grads=[grad],
                check_pir=self.check_pir,
            )

5. 测试结果
   • 本地单元测试: 在修复C++ Kernel并重新编译后，在feature/fix_take_along_axis_0size分支上运行添加的OpTest单元测试，结果为 OK，证明修复成功。
   • PaddleAPITest回测: 由于此API存在参数名不兼容问题，--accuracy模式无法使用。在--paddle_only模式下，修复后可顺利通过。

pcard-67164

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[luotao1]

PR 标题请按照规范

[straigrand]

PR 标题请按照规范

好的，已完成修改。

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 28.57143% with 5 lines in your changes missing coverage. Please review.

Please upload report for BASE (develop@409c59d). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
paddle/phi/kernels/cpu/take_along_axis_kernel.cc	28.57%	5 Missing ⚠️

❌ Your patch status has failed because the patch coverage (28.57%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             develop   #73736   +/-   ##
==========================================
  Coverage           ?   28.57%           
==========================================
  Files              ?        1           
  Lines              ?        7           
  Branches           ?        0           
==========================================
  Hits               ?        2           
  Misses             ?        5           
  Partials           ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[DanielSun11]

“根据任务要求，首先尝试使用PaddleAPITest复现BUG。在PaddleAPITest的--accuracy=True精度对比模式下，由于paddle.take_along_axis(arr, index, axis)与torch.take_along_dim(input, indices, dim)的参数名不统一，反复出现TypeError: missing a required argument的参数绑定错误，无法准确定位问题。
切换到PaddleAPITest的--paddle_only=True模式后，成功触发了Python层的TypeError: take_along_axis() got an unexpected keyword argument 'index'报错，这暴露了API前后端参数名不一致的问题。” 这个问题请去paddleapitest的仓库提个issue，这应该是paddleapitest的bug

[DanielSun11]

LGTM