From 791637cfb50711f439866f7f153b13322e726220 Mon Sep 17 00:00:00 2001
From: Huihuang Zheng <zhhsplendid@163.com>
Date: Tue, 17 Jan 2023 13:19:24 +0800
Subject: [PATCH] Support 0d Tensor in ConditionalBlockOp (#49842)

Support 0d Tensor in ConditionalBlockOp

1. Add dygraph 0d tensor support for ConditionalBlockOp
2. Set scalar loss shape when `append_backward`
---
 python/paddle/fluid/backward.py               |  3 +-
 .../paddle/fluid/tests/unittests/test_cond.py | 32 ++++++++++++++++++-
 python/paddle/static/nn/control_flow.py       |  2 +-
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/python/paddle/fluid/backward.py b/python/paddle/fluid/backward.py
index df975d06a4..19da1ccdff 100755
--- a/python/paddle/fluid/backward.py
+++ b/python/paddle/fluid/backward.py
@@ -385,12 +385,13 @@ def _create_op_desc_(op_type, inputs, outputs, attrs):
 
 
 def _create_loss_op_desc_(loss):
+    create_shape = [] if len(loss.shape) == 0 else [1]
     op_desc = _create_op_desc_(
         "fill_constant",
         {},
         {"Out": [_append_grad_suffix_(loss.name)]},
         {
-            "shape": [1],
+            "shape": create_shape,
             "value": 1.0,
             "dtype": loss.dtype,
             "force_cpu": False,
diff --git a/python/paddle/fluid/tests/unittests/test_cond.py b/python/paddle/fluid/tests/unittests/test_cond.py
index 9769aa8df4..1b1c08e51e 100644
--- a/python/paddle/fluid/tests/unittests/test_cond.py
+++ b/python/paddle/fluid/tests/unittests/test_cond.py
@@ -103,6 +103,7 @@ class TestCondInputOutput(unittest.TestCase):
         exe = fluid.Executor(place)
         (ret,) = exe.run(main_program, fetch_list=[out.name])
         np.testing.assert_allclose(np.asarray(ret), np.array(2), rtol=1e-05)
+        self.assertEqual(ret.shape, ())
 
     def test_0d_tensor_as_cond(self):
         """
@@ -129,7 +130,7 @@ class TestCondInputOutput(unittest.TestCase):
             y = paddle.full(shape=[], dtype='float32', fill_value=0.23)
             pred = paddle.greater_equal(y, x)
             out = paddle.static.nn.cond(pred, true_func, false_func)
-            # out is one tensor
+            # out is a tensor
 
         place = (
             fluid.CUDAPlace(0)
@@ -168,14 +169,41 @@ class TestCondInputOutput(unittest.TestCase):
             if core.is_compiled_with_cuda()
             else fluid.CPUPlace()
         )
+
         exe = fluid.Executor(place)
         ret = exe.run(main_program, fetch_list=[out.name, a.grad_name])
         np.testing.assert_allclose(
             np.asarray(ret[0]), np.array(2.0), rtol=1e-05
         )
+        self.assertEqual(ret[0].shape, ())
         np.testing.assert_allclose(
             np.asarray(ret[1]), np.array(-1.0), rtol=1e-05
         )
+        self.assertEqual(ret[1].shape, ())
+
+    def test_0d_tensor_dygraph(self):
+        """
+        pseudocode:
+
+        a = -2.0
+        if a >= 0:
+            return a
+        else:
+            return -a
+        """
+        paddle.disable_static()
+        a = paddle.full(shape=[], dtype='float32', fill_value=-2.0)
+        a.stop_gradient = False
+        out = paddle.static.nn.cond(a >= 0, lambda: a, lambda: -a)
+        out.backward()
+
+        np.testing.assert_allclose(np.asarray(out), np.array(2.0), rtol=1e-05)
+        self.assertEqual(out.shape, [])
+
+        np.testing.assert_allclose(
+            np.asarray(a.grad), np.array(-1.0), rtol=1e-05
+        )
+        self.assertEqual(a.grad.shape, [])
 
     def test_return_var_tuple(self):
         """
@@ -527,9 +555,11 @@ class TestCondNestedControlFlow(unittest.TestCase):
         np.testing.assert_allclose(
             np.asarray(ret[0]), np.array(7.0), rtol=1e-05
         )
+        self.assertEqual(ret[0].shape, ())
         np.testing.assert_allclose(
             np.asarray(ret[1]), np.array(2.0), rtol=1e-05
         )
+        self.assertEqual(ret[1].shape, ())
 
     def test_cond_op_in_condition(self):
         paddle.enable_static()
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
index d21d95b097..03381b424a 100644
--- a/python/paddle/static/nn/control_flow.py
+++ b/python/paddle/static/nn/control_flow.py
@@ -969,7 +969,7 @@ def cond(pred, true_fn=None, false_fn=None, name=None, return_names=None):
     if _non_static_mode():
         assert isinstance(pred, Variable), "The pred in cond must be Variable"
         assert pred.size == 1, "condition input's numel should be 1"
-        pred = pred.numpy()[0]
+        pred = pred.numpy().item()
         if pred:
             if true_fn is not None:
                 if not callable(true_fn):
-- 
GitLab