[0-size Tensor Retest No.4] fix shape diff for matrix_rank when input is 0-size tensor (#73263)

DanielSun11 · web-flow · commit 3a63fd56c4d5 · 2025-06-12T19:08:24.000+08:00
* fix the matrix rank shape diff  [0-size input]

* rm unused func
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -1361,7 +1361,16 @@ bool MatrixRankTolOpInferSymbolicShape(
                     common::errors::InvalidArgument(
                         "The dims of input must be greater than 2"));
   bool hermitian = GetBoolAttr(op, "hermitian");
-  if (hermitian) {
+  const auto &GetProduct = [&](const auto &dim_exprs) {
+    symbol::DimExpr product{1};
+    for (const auto &dim_expr : dim_exprs) {
+      product = product * dim_expr;
+    }
+    return product;
+  };
+  const auto &x_numel = GetProduct(x_shape);
+
+  if (hermitian && x_numel != 0) {
     infer_context->AddEqualCstr(x_shape[x_rank - 2], x_shape[x_rank - 1]);
   }
   std::vector<symbol::DimExpr> x_shape_batch = [&] {
@@ -1385,6 +1394,12 @@ bool MatrixRankTolOpInferSymbolicShape(
 
   const std::vector<symbol::DimExpr> shapes = [&] {
     std::vector<symbol::DimExpr> shapes;
+    if (x_numel == 0) {
+      if (x_rank == 2)
+        return shapes;  // return empty shape
+      else
+        return x_shape_batch;  // return x_shape[:-2]
+    }
     symbol::DimExprBuilder builder;
     for (size_t i = 0; i < x_shape_batch.size(); i++) {
       if (x_shape_batch[i] == tol_shape[i]) {
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
@@ -2168,6 +2168,15 @@ bool MatrixRankOpInferSymbolicShape(
       infer_context->GetShapeOrDataForValue(op->operand_source(0));
   const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
 
+  const auto &GetProduct = [&](const auto &dim_exprs) {
+    symbol::DimExpr product{1};
+    for (const auto &dim_expr : dim_exprs) {
+      product = product * dim_expr;
+    }
+    return product;
+  };
+  const auto &x_numel = GetProduct(x_shape);
+
   // 确保输入x的维度大于等于2
   PADDLE_ENFORCE_GE(x_shape.size(),
                     2,
@@ -2177,8 +2186,8 @@ bool MatrixRankOpInferSymbolicShape(
   // 获取Hermitian属性
   bool hermitian = op->attribute<pir::BoolAttribute>("hermitian").data();
 
-  // 如果hermitian为true，确保输入x是方阵
-  if (hermitian) {
+  // 如果hermitian为true，确保输入x是方阵,0-size Tensor不需要此检查
+  if (hermitian && x_numel != 0) {
     infer_context->AddEqualCstr(x_shape[x_shape.size() - 2],
                                 x_shape[x_shape.size() - 1]);
   }
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -3177,17 +3177,24 @@ void MatrixRankTolInferMeta(const MetaTensor& x,
                     common::errors::InvalidArgument(
                         "The dims of input must be greater than 2"));
 
-  if (hermitian) {
+  if (hermitian && x.numel() != 0) {
     int64_t rows = static_cast<int64_t>(dim_x[dim_x.size() - 2]);
     int64_t cols = static_cast<int64_t>(dim_x[dim_x.size() - 1]);
+    // if x is 0 size tensor,ignore rows == cols check.
     PADDLE_ENFORCE_EQ(rows,
                       cols,
                       common::errors::InvalidArgument(
                           "if hermitian == true, matrix should be n*n"));
   }
   DDim dim_x_batch = detail::CheckAndGetOutputDim(dim_x);
   auto dim_tol = atol_tensor.dims();
-  if (dim_x_batch == dim_tol) {
+  if (x.numel() == 0) {
+    if (dim_x.size() == 2) {
+      out->set_dims(common::make_ddim({}));
+    } else {
+      out->set_dims(dim_x_batch);
+    }
+  } else if (dim_x_batch == dim_tol) {
     out->set_dims(dim_x_batch);
   } else {
     int max_dim = std::max(dim_x_batch.size(), dim_tol.size());
diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -26,6 +26,17 @@ limitations under the License. */
 #include "paddle/phi/kernels/impl/box_coder.h"
 
 namespace phi {
+namespace detail {
+// Used in MatrixRankAtolRtolInferMeta
+static DDim CheckAndGetOutputDim(const DDim& dim_x) {
+  auto x_vec = common::vectorize(dim_x);
+  if (x_vec.size() == 2) {
+    return common::make_ddim({});
+  }
+  x_vec.erase(x_vec.end() - 2, x_vec.end());
+  return common::make_ddim(x_vec);
+}
+}  // namespace detail
 
 void AccuracyInferMeta(const MetaTensor& out,
                        const MetaTensor& indice,
@@ -1493,6 +1504,18 @@ void MatrixRankAtolRtolInferMeta(const MetaTensor& x,
                                  const MetaTensor& rtol,
                                  bool hermitian,
                                  MetaTensor* out) {
+  if (x.numel() == 0) {
+    auto dim_x = x.dims();
+    PADDLE_ENFORCE_GE(dim_x.size(),
+                      2,
+                      common::errors::InvalidArgument(
+                          "The dims of input must be greater than 2"));
+
+    DDim dim_x_batch = detail::CheckAndGetOutputDim(dim_x);
+    out->set_dims(dim_x_batch);
+    out->share_lod(x);
+    return;
+  }
   MatrixRankTolInferMeta(x, atol, true, hermitian, out);
 }
 
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -2534,9 +2534,10 @@ void MatrixRankInferMeta(const MetaTensor& x,
                     common::errors::InvalidArgument(
                         "The dims of input must be greater than 2."));
 
-  if (hermitian) {
+  if (hermitian && x.numel() != 0) {
     int rows = static_cast<int>(dim_x[dim_x.size() - 2]);
     int cols = static_cast<int>(dim_x[dim_x.size() - 1]);
+    // if x is 0-size Tensor,ignore rows == cols check.
     PADDLE_ENFORCE_EQ(rows,
                       cols,
                       common::errors::InvalidArgument(
diff --git a/test/legacy_test/test_matrix_rank_op.py b/test/legacy_test/test_matrix_rank_op.py
@@ -474,6 +474,21 @@ def _test_matrix_rank_dynamic_gpu(self):
             )
             np.testing.assert_allclose(y_valid2, y_valid2_real, rtol=1e-05)
 
+            x_valid3 = paddle.full((2, 0, 7, 7), 1.0, dtype='float64')
+            tol = paddle.full((2, 0), 1.0, dtype='float32')
+            y_valid3 = paddle.linalg.matrix_rank(x_valid3, tol)
+            self.assertEqual(y_valid3.shape, x_valid3.shape[:-2])
+
+            x_valid4 = paddle.full((2, 0, 7, 7), 1.0, dtype='float64')
+            atol = paddle.full((2, 0), 1.0, dtype='float32')
+            y_valid4 = paddle.linalg.matrix_rank(x_valid4, atol=atol)
+            self.assertEqual(y_valid4.shape, x_valid4.shape[:-2])
+
+            x_valid5 = paddle.full((0, 7), 1.0, dtype='float64')
+            tol = paddle.full((1), 1.0, dtype='float32')
+            y_valid5 = paddle.linalg.matrix_rank(x_valid5, tol)
+            self.assertEqual(y_valid5.shape, y_valid5.shape[:-2])
+
     def test_matrix_rank_tensor(self):
         for place in self._get_places():
             self._test_matrix_rank_static(place)
