[0-size Tensor No.141] Add 0-size Tensor support for nanmedian

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc

@@ -2314,15 +2314,40 @@ bool NanmedianOpInferSymbolicShape(
   if (mode == "avg") {
     median_shape.emplace_back(2);
   }
-  infer_context->SetShapeOrDataForValue(
-      op->result(0),
-      symbol::ShapeOrDataDimExprs{
-          symbol::TensorShapeOrDataDimExprs(out_shape)});
-  infer_context->SetShapeOrDataForValue(
-      op->result(1),
-      symbol::ShapeOrDataDimExprs{
-          symbol::TensorShapeOrDataDimExprs(median_shape)});
 
+  const auto &IsZero = [&](const symbol::DimExpr &dim_expr) {
+    if (dim_expr.isa<int64_t>()) {
+      return dim_expr.dyn_cast<int64_t>() == static_cast<int64_t>(0);
+    }
+    return false;
+  };
+  bool size_0 = false;
+  for (int i = 0; i < x_shape.size(); i++) {
+    if (IsZero(x_shape.at(i))) {
+      size_0 = true;
+      break;
+    }
+  }
+  if (size_0) {
+    std::vector<symbol::DimExpr> x_numel_0_shape = {};
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_numel_0_shape)});
+    infer_context->SetShapeOrDataForValue(
+        op->result(1),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_numel_0_shape)});
+  } else {
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(out_shape)});
+    infer_context->SetShapeOrDataForValue(
+        op->result(1),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(median_shape)});
+  }
   return true;
 }
 

paddle/phi/infermeta/unary.cc

@@ -2854,6 +2854,11 @@ void NanmedianInferMeta(const MetaTensor& x,
   }
   median_index->set_dtype(DataType::INT64);
   median_index->set_dims(make_ddim(median_dim));
+
+  if (x.numel() == 0) {
+    out->set_dims(make_ddim({}));
+    median_index->set_dims(make_ddim({}));
+  }
 }
 
 void NMSInferMeta(const MetaTensor& x, float threshold, MetaTensor* out) {

paddle/phi/kernels/cpu/nanmedian_grad_kernel.cc

@@ -96,6 +96,10 @@ void NanmedianGradKernel(const Context& dev_ctx,
                          bool keepdim UNUSED,
                          const std::string& mode,
                          DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   DenseTensor tmp_x;
   auto rank = x.dims().size();
   if ((axes.size() == 0) || rank <= 1) {

paddle/phi/kernels/cpu/nanmedian_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/nanmedian_utils.h"
 #include "paddle/phi/kernels/top_k_kernel.h"
 
@@ -218,6 +219,16 @@ void NanmedianKernel(const Context& dev_ctx,
                      const std::string& mode,
                      DenseTensor* out,
                      DenseTensor* median_index) {
+  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(median_index->dims())),
+        0,
+        median_index);
+    return;
+  }
   DenseTensor tmp_x;
   auto rank = x.dims().size();
   if ((axes.size() == 0) || rank <= 1) {

paddle/phi/kernels/gpu/nanmedian_grad_kernel.cu

@@ -111,6 +111,10 @@ void NanmedianGradKernel(const Context& dev_ctx,
                          bool keepdim UNUSED,
                          const std::string& mode,
                          DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   DenseTensor tmp_x;
   auto rank = x.dims().size();
   if ((axes.size() == 0) || rank <= 1) {

paddle/phi/kernels/gpu/nanmedian_kernel.cu

@@ -356,6 +356,16 @@ void NanmedianKernel(const Context& dev_ctx,
                      const std::string& mode,
                      DenseTensor* out,
                      DenseTensor* median_index) {
+  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(median_index->dims())),
+        0,
+        median_index);
+    return;
+  }
   DenseTensor tmp_x;
   auto rank = x.dims().size();
   if ((axes.size() == 0) || rank <= 1) {

test/legacy_test/test_nanmedian.py

@@ -617,6 +617,23 @@ def test_check_grad(self):
         self.check_grad(['X'], 'Out', check_pir=True)
 
 
+class TestNanmedianOp_ZeroSize(TestNanmedianFP16Op):
+    def setUp(self):
+        self.op_type = "nanmedian"
+        self.python_api = paddle.nanmedian
+        self.public_python_api = paddle.nanmedian
+        self.dtype = np.float64
+        self.python_out_sig = ["Out"]
+        X = np.random.random((100, 0)).astype('float64')
+        Out = np.nanmedian(X)
+        indices = np.zeros_like(Out, dtype='int64')
+        self.inputs = {'X': X}
+        self.outputs = {'Out': Out, 'MedianIndex': indices}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, equal_nan=True)
+
+
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
     or not core.is_bfloat16_supported(core.CUDAPlace(0)),