[0-size Tensor Job2 No.48] Add 0-size Tensor support for paddle.median (#73477)

co63oc · web-flow · commit aaf077d4c57b · 2025-06-24T10:07:03.000+08:00
* Fix

* Fix
diff --git a/paddle/phi/kernels/cpu/top_k_kernel.cc b/paddle/phi/kernels/cpu/top_k_kernel.cc
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
@@ -159,19 +160,26 @@ void TopkKernel(const Context& dev_ctx,
   }
 
   int k = k_scalar.to<int>();
-  PADDLE_ENFORCE_GE(
-      x.numel(),
-      k,
-      errors::InvalidArgument(
-          "x has only %d element, can not find %d top values.", x.numel(), k));
-
+  // out shape [-1]
   if (k_scalar.FromTensor()) {
     auto out_dims = out->dims();
     // according to axis to set K value in the dim
     out_dims[axis] = k;
     out->Resize(out_dims);
     indices->Resize(out_dims);
   }
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
+    phi::Full<int64_t, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(indices->dims())), 0, indices);
+    return;
+  }
+  PADDLE_ENFORCE_GE(
+      x.numel(),
+      k,
+      errors::InvalidArgument(
+          "x has only %d element, can not find %d top values.", x.numel(), k));
 
   T* out_data = dev_ctx.template Alloc<T>(out);
   int64_t* indices_data = dev_ctx.template Alloc<int64_t>(indices);
diff --git a/paddle/phi/kernels/gpu/top_k_kernel.cu b/paddle/phi/kernels/gpu/top_k_kernel.cu
@@ -20,10 +20,10 @@
 #include "paddle/phi/common/bfloat16.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/gather.cu.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/top_k_function_cuda.h"
-
 namespace phi {
 
 #define FIXED_BLOCK_DIM_BASE(dim, ...) \
@@ -83,17 +83,25 @@ void TopkKernel(const Context& dev_ctx,
   if (axis < 0) axis += in_dims.size();
 
   int k = k_scalar.to<int>();
-  PADDLE_ENFORCE_GE(
-      x.numel(),
-      k,
-      errors::InvalidArgument(
-          "x has only %d element, can not find %d top values.", x.numel(), k));
+  // out shape [-1]
   if (k_scalar.FromTensor()) {
     phi::DDim out_dims = out->dims();
     out_dims[axis] = k;
     out->Resize(out_dims);
     indices->Resize(out_dims);
   }
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
+    phi::Full<int64_t, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(indices->dims())), 0, indices);
+    return;
+  }
+  PADDLE_ENFORCE_GE(
+      x.numel(),
+      k,
+      errors::InvalidArgument(
+          "x has only %d element, can not find %d top values.", x.numel(), k));
 
   const auto& out_dims = out->dims();
 
diff --git a/paddle/phi/kernels/xpu/top_k_kernel.cc b/paddle/phi/kernels/xpu/top_k_kernel.cc
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/xpu/xpu_mem_util.h"
 namespace phi {
@@ -52,19 +53,26 @@ void TopkKernel(const Context& dev_ctx,
   }
 
   int64_t k = k_scalar.to<int64_t>();
-  PADDLE_ENFORCE_GE(
-      x.numel(),
-      k,
-      errors::InvalidArgument(
-          "x has only %d element, can not find %d top values.", x.numel(), k));
-
+  // out shape [-1]
   if (k_scalar.FromTensor()) {
     auto out_dims_ = out->dims();
     // according to axis to set K value in the dim
     out_dims_[axis] = k;
     out->Resize(out_dims_);
     indices->Resize(out_dims_);
   }
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
+    phi::Full<int64_t, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(indices->dims())), 0, indices);
+    return;
+  }
+  PADDLE_ENFORCE_GE(
+      x.numel(),
+      k,
+      errors::InvalidArgument(
+          "x has only %d element, can not find %d top values.", x.numel(), k));
 
   const T* in_data = x.data<T>();
   int64_t* indices_data = dev_ctx.template Alloc<int64_t>(indices);
diff --git a/python/paddle/tensor/stat.py b/python/paddle/tensor/stat.py
@@ -598,10 +598,6 @@ def median(
     if not isinstance(x, (Variable, paddle.pir.Value)):
         raise TypeError("In median, the input x should be a Tensor.")
 
-    if in_dynamic_mode() and x.size == 0:
-        # TODO: Currently, `__eq__` don't support arguments (`pir.Value` & `int`)
-        raise ValueError("In median, the size of input x should not be 0.")
-
     is_flatten = False
     dims = len(x.shape)
     if dims == 0:
@@ -658,7 +654,7 @@ def median(
             keepdim=True,
         )
     else:  # mode == 'min'
-        if sz & 1 == 0:
+        if sz & 1 == 0 and kth != 0:
             out_tensor = paddle.slice(
                 tensor_topk, axes=[axis], starts=[kth - 1], ends=[kth]
             )
diff --git a/test/legacy_test/test_median.py b/test/legacy_test/test_median.py
@@ -50,7 +50,7 @@ def np_medain_min(data, keepdims=False):
     return np_res + np.sum(np.isnan(data).astype(data.dtype) * data)
 
 
-def np_medain_min_axis(data, axis=None, keepdims=False):
+def np_median_min_axis(data, axis=None, keepdims=False):
     data = copy.deepcopy(data)
     if axis is None:
         return np_medain_min(data, keepdims)
@@ -232,7 +232,7 @@ class TestMedianMin(unittest.TestCase):
     def static_single_test_median(self, lis_test):
         paddle.enable_static()
         x, axis, keepdims = lis_test
-        res_np = np_medain_min_axis(x, axis=axis, keepdims=keepdims)
+        res_np = np_median_min_axis(x, axis=axis, keepdims=keepdims)
         main_program = paddle.static.Program()
         startup_program = paddle.static.Program()
         exe = paddle.static.Executor()
@@ -245,7 +245,7 @@ def static_single_test_median(self, lis_test):
 
     def dygraph_single_test_median(self, lis_test):
         x, axis, keepdims = lis_test
-        res_np = np_medain_min_axis(x, axis=axis, keepdims=keepdims)
+        res_np = np_median_min_axis(x, axis=axis, keepdims=keepdims)
         if axis is None:
             res_pd = paddle.median(
                 paddle.to_tensor(x), axis, keepdims, mode='min'
@@ -335,7 +335,7 @@ def test_float16(self):
             for keepdims in [False, True]
         ]
         for axis, keepdims in lis_tests:
-            res_np = np_medain_min_axis(x, axis=axis, keepdims=keepdims)
+            res_np = np_median_min_axis(x, axis=axis, keepdims=keepdims)
             if axis is None:
                 res_pd = paddle.median(
                     paddle.to_tensor(x), axis, keepdims, mode='min'
@@ -357,5 +357,49 @@ def test_output_dtype(self):
             np.testing.assert_equal(res.numpy().dtype, np.dtype(inp_dtype))
 
 
+class TestMedianAvg_ZeroSize(unittest.TestCase):
+    def dygraph_single_test_median(self, lis_test):
+        x, axis, keepdims = lis_test
+        res_np = np.median(x, axis=axis, keepdims=keepdims)
+        x_pd = paddle.to_tensor(x)
+        x_pd.stop_gradient = False
+        res_pd = paddle.median(x_pd, axis, keepdims)
+        np.testing.assert_allclose(res_pd.numpy(), res_np)
+        paddle.sum(res_pd).backward()
+        np.testing.assert_allclose(x_pd.grad.shape, x_pd.shape)
+
+    def test_median_dygraph(self):
+        paddle.disable_static()
+        h = 0
+        w = 4
+        l = 2
+        x = np.arange(h * w * l).reshape([h, w, l])
+        self.dygraph_single_test_median([x, 1, False])
+
+
+class TestMedianMin_ZeroSize(unittest.TestCase):
+
+    def dygraph_single_test_median(self, lis_test):
+        x, axis, keepdims = lis_test
+        res_np = np_median_min_axis(x, axis=axis, keepdims=keepdims)
+        x_pd = paddle.to_tensor(x)
+        x_pd.stop_gradient = False
+        if axis is None:
+            res_pd = paddle.median(x_pd, axis, keepdims, mode='min')
+        else:
+            res_pd, _ = paddle.median(x_pd, axis, keepdims, mode='min')
+        np.testing.assert_allclose(res_pd.numpy(), res_np)
+        paddle.sum(res_pd).backward()
+        np.testing.assert_allclose(x_pd.grad.shape, x_pd.shape)
+
+    def test_median_dygraph(self):
+        paddle.disable_static()
+        h = 0
+        w = 4
+        l = 2
+        x = np.arange(h * w * l).reshape([h, w, l]).astype("float32")
+        self.dygraph_single_test_median([x, 1, False])
+
+
 if __name__ == '__main__':
     unittest.main()
