[0-size Tensor No.11、262、287] Add 0-size Tensor support for argsort/sort API. (#72872)

enkilee · web-flow · commit cf56cbd9de3a · 2025-05-30T14:50:39.000+08:00
* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix
diff --git a/paddle/phi/kernels/xpu/argsort_kernel.cc b/paddle/phi/kernels/xpu/argsort_kernel.cc
@@ -17,6 +17,7 @@
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/backends/xpu/xpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
@@ -99,6 +100,15 @@ void ArgsortKernel(const Context& dev_ctx,
                    DenseTensor* indices) {
   auto in_dims = input.dims();
   auto rank = in_dims.size();
+
+  if (input.numel() == 0) {
+    output->Resize(in_dims);
+    indices->Resize(in_dims);
+    dev_ctx.template Alloc<T>(output);
+    dev_ctx.template Alloc<int64_t>(indices);
+    return;
+  }
+
   axis = (axis < 0) ? (in_dims.size() + axis) : axis;
   int64_t n = in_dims[axis];
 
diff --git a/test/xpu/test_argsort_op_xpu.py b/test/xpu/test_argsort_op_xpu.py
@@ -96,6 +96,65 @@ def test_check_grad(self):
             self.check_grad_with_place(self.place, {'X'}, 'Out')
 
 
+class XPUTestArgsortOp_0D(XPUOpTestWrapper):
+    def __init__(self):
+        self.op_name = 'argsort'
+        self.use_dynamic_create_class = False
+
+        class TestArgsortOpCase1(XPUOpTest):
+            def setUp(self):
+                self.set_xpu()
+                self.op_type = "argsort"
+                self.place = paddle.XPUPlace(0)
+                self.dtype = self.in_type
+                self.input_shape = 0
+                self.axis = (
+                    -1 if not hasattr(self, 'init_axis') else self.init_axis
+                )
+                self.descending = (
+                    False
+                    if not hasattr(self, 'init_descending')
+                    else self.init_descending
+                )
+
+                if self.dtype == np.float32:
+                    self.x = np.random.random(self.input_shape).astype(
+                        self.dtype
+                    )
+                else:
+                    self.x = np.random.choice(
+                        low=-1000, high=1000, size=self.input_shape
+                    ).astype(self.dtype)
+
+            self.inputs = {"X": self.x}
+            self.attrs = {"axis": self.axis, "descending": self.descending}
+            self.get_output()
+            self.outputs = {"Out": self.sorted_x, "Indices": self.indices}
+
+    def get_output(self):
+        if self.descending:
+            self.indices = np.flip(
+                np.argsort(self.x, kind='heapsort', axis=self.axis),
+                self.axis,
+            )
+            self.sorted_x = np.flip(
+                np.sort(self.x, kind='heapsort', axis=self.axis), self.axis
+            )
+        else:
+            self.indices = np.argsort(self.x, kind='heapsort', axis=self.axis)
+            self.sorted_x = np.sort(self.x, kind='heapsort', axis=self.axis)
+
+    def set_xpu(self):
+        self.__class__.use_xpu = True
+        self.__class__.no_need_check_grad = True
+
+    def test_check_output(self):
+        self.check_output_with_place(self.place)
+
+    def test_check_grad(self):
+        self.check_grad_with_place(self.place, {'X'}, 'Out')
+
+
 class XPUTestArgsortOp_LargeN(XPUOpTestWrapper):
     def __init__(self):
         self.op_name = 'argsort'
