[0D-Tensor] CINN supports squeeze, fix infershape and GetPositiveAxes (#55333)

jiahy0825 · web-flow · commit bb0df46836ff · 2023-07-12T19:40:55.000+08:00
diff --git a/paddle/cinn/hlir/op/elementwise.cc b/paddle/cinn/hlir/op/elementwise.cc
@@ -705,10 +705,6 @@ std::vector<std::vector<int>> InferShapeForSqueeze(
 
   VLOG(4) << "The output calculated in Squeeze: "
           << cinn::utils::Join(output_shape, ", ");
-
-  if (output_shape.size() == 0) {
-    output_shape.push_back(1);
-  }
   return {output_shape};
 }
 
diff --git a/paddle/cinn/utils/functional.cc b/paddle/cinn/utils/functional.cc
@@ -23,7 +23,7 @@ std::vector<int> GetPositiveAxes(const std::vector<int>& axes, int rank) {
   std::vector<int> new_axes(axes.size());
   for (int i = 0; i < axes.size(); ++i) {
     int axis = axes[i] + (axes[i] < 0 ? rank : 0);
-    CHECK(axis >= 0 && axis < rank)
+    CHECK(axis >= 0 && (rank == 0 || axis < rank))
         << "The axis should in [" << -rank << ", " << rank << "), but axes["
         << i << "]=" << axes[i] << " not.";
     new_axes[i] = axis;
diff --git a/test/cinn/ops/test_zero_dim_tensor.py b/test/cinn/ops/test_zero_dim_tensor.py
@@ -713,5 +713,62 @@ def test_check_results(self):
         self.check_outputs_and_grads()
 
 
+@OpTestTool.skip_if(
+    not is_compiled_with_cuda(), "x86 test will be skipped due to timeout."
+)
+class TestSqueezeOp(OpTest):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.init_input()
+
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, []).astype(self.dtype),
+        }
+        self.squeeze_axex = [0]
+        self.target_shape = ()
+
+    def build_paddle_program(self, target):
+        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
+        out = paddle.squeeze(x, axis=self.squeeze_axex)
+
+        self.paddle_outputs = [out]
+
+    def build_cinn_program(self, target):
+        builder = NetBuilder("squeeze_op")
+        x = builder.create_input(
+            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x"
+        )
+        out = builder.squeeze(x, self.squeeze_axex)
+
+        prog = builder.build()
+        res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]], [out])
+
+        self.cinn_outputs = res
+        self.assertEqual(res[0].shape, self.target_shape)
+
+    def test_check_results(self):
+        self.check_outputs_and_grads()
+
+
+class TestSqueezeOp1D(TestSqueezeOp):
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, [1]).astype(self.dtype),
+        }
+        self.squeeze_axex = []
+        self.target_shape = ()
+
+
+class TestSqueezeOp2D(TestSqueezeOp):
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, [1, 1]).astype(self.dtype),
+        }
+        self.squeeze_axex = [0, 1]
+        self.target_shape = ()
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -705,10 +705,6 @@ std::vector<std::vector<int>> InferShapeForSqueeze(`
`705`	`705`
`706`	`706`	`VLOG(4) << "The output calculated in Squeeze: "`
`707`	`707`	`<< cinn::utils::Join(output_shape, ", ");`
`708`		`-`
`709`		`- if (output_shape.size() == 0) {`
`710`		`- output_shape.push_back(1);`
`711`		`- }`
`712`	`708`	`return {output_shape};`
`713`	`709`	`}`
`714`	`710`