[Zero-Dim] support 0D Tensor for while_loop op

paddle/fluid/operators/controlflow/while_op.cc

@@ -100,12 +100,12 @@ class WhileOp : public framework::OperatorBase {
 
     auto &cond = scope.FindVar(Input(kCondition))->Get<phi::DenseTensor>();
     PADDLE_ENFORCE_EQ(
-        cond.dims(),
-        phi::make_ddim({1}),
+        cond.numel(),
+        1,
         platform::errors::InvalidArgument(
-            "The shape of Input(Condition) of WhileOp must be 1. But now "
-            "the Condition's shape is ",
-            cond.dims().to_str(),
+            "The numel of Input(Condition) of WhileOp must be 1. But now "
+            "the Condition's numel is ",
+            cond.numel(),
             ".\n"));
 
 #ifdef PADDLE_WITH_MKLDNN

python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py

@@ -2297,6 +2297,47 @@ def test_sequence_pad(self):
         res = self.exe.run(prog, feed={"x": x_tensor}, fetch_list=[out])
         self.assertEqual(res[0].shape, (3, 4, 2))
 
+    @prog_scope()
+    def test_while_loop(self):
+        def cond(i, x):
+            return paddle.less_than(i, eleven)
+
+        def body(i, x):
+            x = paddle.multiply(x=i, y=i)
+            i = paddle.increment(i)
+            return [i, x]
+
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, paddle.static.Program()):
+            i = paddle.static.data(name='i', shape=[], dtype='float32')
+            i.stop_gradient = False
+            eleven = paddle.full([], 11, 'float32')
+            one = paddle.full([], 1, 'float32')
+            x = paddle.static.data(name='x', shape=[], dtype='float32')
+            x.stop_gradient = False
+
+            out_i, out_x = paddle.static.nn.while_loop(cond, body, [i, x])
+            paddle.static.append_backward(out_x.sum())
+
+        place = (
+            paddle.CUDAPlace(0)
+            if paddle.device.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+        feed_i = np.ones([]).astype('float32')
+        feed_x = np.ones([]).astype('float32')
+
+        res = self.exe.run(
+            main_program,
+            feed={'i': feed_i, 'x': feed_x},
+            fetch_list=[out_i.name, out_x.name, i.grad_name, x.grad_name],
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[3].shape, ())
+
 
 # Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.
 class TestNoBackwardAPI(unittest.TestCase):