[0-size Tensor Retest No.5] fix output shape diff for paddle.nanmedian (#73267)

DanielSun11 · web-flow · commit a6fcb93b5eb0 · 2025-06-12T19:11:24.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -2324,39 +2324,14 @@ bool NanmedianOpInferSymbolicShape(
     median_shape.emplace_back(2);
   }
 
-  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 (size_t 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)});
-  }
+  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;
 }
 
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -2885,11 +2885,6 @@ 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) {
diff --git a/test/legacy_test/test_nanmedian.py b/test/legacy_test/test_nanmedian.py
@@ -617,21 +617,26 @@ 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}
+class TestNanmedianZeroSize(unittest.TestCase):
+    def init_data(self):
+        self.x_shape = [100, 0]
+        self.expect_out = np.array(np.nan, dtype='float32')
+        self.axis = None
+
+    def test_zero(self):
+        self.init_data()
+        x = paddle.randn(self.x_shape)
+        out = paddle.nanmedian(x, axis=self.axis)
+        np.testing.assert_allclose(
+            out.numpy(), self.expect_out, rtol=1e-05, equal_nan=True
+        )
 
-    def test_check_output(self):
-        self.check_output(check_pir=True, equal_nan=True)
+
+class TestNanmedianZeroSize1(TestNanmedianZeroSize):
+    def init_data(self):
+        self.x_shape = [100, 0, 100]
+        self.expect_out = np.zeros(self.x_shape[:-1], dtype='float32')
+        self.axis = -1
 
 
 @unittest.skipIf(

Original file line number	Diff line number	Diff line change
`@@ -2885,11 +2885,6 @@ void NanmedianInferMeta(const MetaTensor& x,`
`2885`	`2885`	`}`
`2886`	`2886`	`median_index->set_dtype(DataType::INT64);`
`2887`	`2887`	`median_index->set_dims(make_ddim(median_dim));`
`2888`		`-`
`2889`		`- if (x.numel() == 0) {`
`2890`		`- out->set_dims(make_ddim({}));`
`2891`		`- median_index->set_dims(make_ddim({}));`
`2892`		`- }`
`2893`	`2888`	`}`
`2894`	`2889`
`2895`	`2890`	`void NMSInferMeta(const MetaTensor& x, float threshold, MetaTensor* out) {`