【Paddle Tensor 规范化第二期】paddle.argsort 支持 0-size tensor (#71094)

fangfangssj · web-flow · commit 3967f448bc55 · 2025-02-13T17:02:08.000+08:00
* support argsort

* fix test

* add test
diff --git a/paddle/phi/kernels/cpu/argsort_kernel.cc b/paddle/phi/kernels/cpu/argsort_kernel.cc
@@ -95,6 +95,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;
   T* out_data = dev_ctx.template Alloc<T>(output);
 
diff --git a/paddle/phi/kernels/gpu/argsort_kernel.cu b/paddle/phi/kernels/gpu/argsort_kernel.cu
@@ -235,6 +235,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;
   const T* in_data = input.data<T>();
   auto size = input.numel();
diff --git a/test/legacy_test/test_argsort_op.py b/test/legacy_test/test_argsort_op.py
@@ -155,6 +155,30 @@ def init(self):
         self.axis = 1
 
 
+class TestArgsortZeroSize(TestArgsort):
+    def init(self):
+        self.input_shape = [0, 3]
+        self.axis = 0
+
+
+class TestArgsortZeroSize2(TestArgsort):
+    def init(self):
+        self.input_shape = [2, 0, 4]
+        self.axis = 0
+
+
+class TestArgsortZeroSize3(TestArgsort):
+    def init(self):
+        self.input_shape = [0, 3]
+        self.axis = 1
+
+
+class TestArgsortZeroSize4(TestArgsort):
+    def init(self):
+        self.input_shape = [2, 0, 4]
+        self.axis = 1
+
+
 class TestStableArgsort(unittest.TestCase):
     def init(self):
         self.input_shape = [
@@ -283,6 +307,34 @@ def init(self):
         )
 
 
+class TestStableArgsortZeroSize(TestStableArgsort):
+    def init(self):
+        self.input_shape = [0, 30]
+        self.data = np.random.rand(*self.input_shape)
+        self.axis = 0
+
+
+class TestStableArgsortZeroSize2(TestStableArgsort):
+    def init(self):
+        self.input_shape = [2, 0, 40]
+        self.data = np.random.rand(*self.input_shape)
+        self.axis = 0
+
+
+class TestStableArgsortZeroSize3(TestStableArgsort):
+    def init(self):
+        self.input_shape = [0, 30]
+        self.data = np.random.rand(*self.input_shape)
+        self.axis = 1
+
+
+class TestStableArgsortZeroSize4(TestStableArgsort):
+    def init(self):
+        self.input_shape = [2, 0, 40]
+        self.data = np.random.rand(*self.input_shape)
+        self.axis = 1
+
+
 class TestArgsortImperative(unittest.TestCase):
     def init(self):
         self.input_shape = [
