【Hackathon 6th No.9】Update ValidateShape to support zero sized tensor flatten -part

paddle/phi/infermeta/unary.cc

@@ -2090,7 +2090,15 @@ static phi::DDim ValidateShape(const std::vector<int64_t> shape,
       output_shape[i] = shape[i];
     } else if (shape[i] == 0) {
       if (static_cast<int>(i) < in_dims.size()) {
-        output_shape[i] = in_dims[static_cast<int>(i)];
+        if (in_size == 0) {
+          // such as [3, 2, 0] -> [0, 0] is [0, 0]; [3, 2, 0] -> [10, 0] is [10,
+          // 0]
+          output_shape[i] = 0;
+        } else {
+          // such as [3, 2, 1] -> [0, 0] is [3, 2]; [3, 2, 1] -> [3, 2, 0] is
+          // [3, 2, 1]
+          output_shape[i] = in_dims[static_cast<int>(i)];
+        }
       } else {
         PADDLE_ENFORCE_EQ(
             in_size,

test/deprecated/legacy_test/test_flatten_contiguous_range_op.py

@@ -580,5 +580,29 @@ def test_ValueError2():
         self.assertRaises(ValueError, test_ValueError2)
 
 
+class TestFlattenZeroSizedTensorAPI(unittest.TestCase):
+    def test_dygraph(self):
+        paddle.disable_static()
+        data = np.random.randn(2, 3, 0)
+        x = paddle.to_tensor(data)
+        out = paddle.flatten(x)
+        out_np = data.flatten()
+        np.testing.assert_equal(out.numpy(), out_np)
+
+    @test_with_pir_api
+    def test_static(self):
+        paddle.enable_static()
+        data = np.random.randn(2, 3, 0)
+        main_prog = paddle.static.Program()
+        with paddle.static.program_guard(main_prog, paddle.static.Program()):
+            x = paddle.static.data(name="x", shape=[2, 3, 0], dtype='float64')
+            out = paddle.flatten(x)
+
+        exe = paddle.static.Executor(place=paddle.CPUPlace())
+        fetch_out = exe.run(main_prog, feed={"x": data}, fetch_list=[out])[0]
+        out_np = data.flatten()
+        np.testing.assert_equal(fetch_out, out_np)
+
+
 if __name__ == "__main__":
     unittest.main()