[OP] Enlarge max rank for unsqueeze from 6 to 8 (#63556)

HydrogenSulfate · web-flow · commit 3334b554bd7c · 2024-04-18T10:49:02.000+08:00
* enlarge max rank for unsqueeze from 6 to 8

* add unitest for unsqueeze rank=7, 8
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -5478,14 +5478,16 @@ void UnsqueezeInferMeta(const MetaTensor& x,
                         const IntArray& axes,
                         MetaTensor* out,
                         MetaConfig config) {
+#define UNSQUEEZE_MAX_RANK_SUPPORTED 8
   const auto& x_dims = x.dims();
   // Validity Check: input tensor dims (<6).
   PADDLE_ENFORCE_LE(x_dims.size(),
-                    6,
+                    UNSQUEEZE_MAX_RANK_SUPPORTED,
                     phi::errors::InvalidArgument(
                         "Invalid "
                         "dimensions, the rank of Input(X) "
-                        "should be in the range of [1, 6] (Eigen limit)"));
+                        "should be in the range of [1, %d] (Eigen limit)",
+                        UNSQUEEZE_MAX_RANK_SUPPORTED));
   if (!config.is_runtime && axes.FromTensor()) {
     // compile time infershape.  set all elements to -1.
     int output_size = static_cast<int>(x.dims().size() + axes.GetData().size());
@@ -5500,6 +5502,7 @@ void UnsqueezeInferMeta(const MetaTensor& x,
     }
     out->set_dtype(x.dtype());
   }
+#undef UNSQUEEZE_MAX_RANK_SUPPORTED
 }
 
 void UnsqueezeWithXShapeInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/kernels/funcs/unsqueeze.h b/paddle/phi/kernels/funcs/unsqueeze.h
@@ -105,16 +105,17 @@ inline DDim GetOutputSqueezeShape(const std::vector<int> squeeze_dims,
 
 inline DDim GetUnsqueezeShape(const std::vector<int64_t> unsqz_dims,
                               const DDim& in_dims) {
+#define UNSQUEEZE_MAX_RANK_SUPPORTED 8
   int output_rank = in_dims.size() + static_cast<int>(unsqz_dims.size());
   int cur_output_rank = in_dims.size();
   std::vector<int64_t> output_shape(output_rank, 0);
-
   // Validity Check: rank range.
   PADDLE_ENFORCE_LE(
       output_rank,
-      6,
+      UNSQUEEZE_MAX_RANK_SUPPORTED,
       phi::errors::InvalidArgument("The output "
-                                   "tensor's rank should be less than 6."));
+                                   "tensor's rank should be less than %d.",
+                                   UNSQUEEZE_MAX_RANK_SUPPORTED));
 
   for (int axis : unsqz_dims) {
     int cur = axis < 0 ? axis + cur_output_rank + 1 : axis;
@@ -148,7 +149,7 @@ inline DDim GetUnsqueezeShape(const std::vector<int64_t> unsqz_dims,
       output_shape[out_idx] = in_dims[in_idx++];
     }
   }
-
+#undef UNSQUEEZE_MAX_RANK_SUPPORTED
   return common::make_ddim(output_shape);
 }
 
diff --git a/test/legacy_test/test_unsqueeze2_op.py b/test/legacy_test/test_unsqueeze2_op.py
@@ -101,6 +101,22 @@ def init_test_case(self):
         self.new_shape = (10, 1, 1, 2, 5, 1)
 
 
+# Test for output rank=7
+class TestUnsqueezeOp5(TestUnsqueezeOp):
+    def init_test_case(self):
+        self.ori_shape = (10, 2, 5)
+        self.axes = (1, 2, 3, 4)
+        self.new_shape = (10, 1, 1, 1, 1, 2, 5)
+
+
+# Test for output rank=8
+class TestUnsqueezeOp6(TestUnsqueezeOp):
+    def init_test_case(self):
+        self.ori_shape = (10, 2, 5)
+        self.axes = (1, 2, 3, 4, 5)
+        self.new_shape = (10, 1, 1, 1, 1, 1, 2, 5)
+
+
 class TestUnsqueezeOp_ZeroDim1(TestUnsqueezeOp):
     def init_test_case(self):
         self.ori_shape = ()
