Support 0d Tensor in ConditionalBlockOp

paddle/fluid/operators/controlflow/conditional_block_op.cc

@@ -384,8 +384,10 @@ class ConditionalBlockGradOp : public ConditionalOp {
     if (!input_tensor.IsInitialized() || input_tensor.numel() == 0) {
       return;
     }
+    if (!input_tensor.meta().is_scalar) {
+      outside_tensor->Resize(input_tensor.dims());
+    }
     VLOG(4) << "Assigning zero to " << outside_tensor;
-    outside_tensor->Resize(input_tensor.dims());
     outside_tensor->mutable_data(place, input_tensor.dtype());
     const platform::DeviceContext *dev_ctx =
         platform::DeviceContextPool::Instance().Get(place);

paddle/phi/core/tensor_utils.cc

@@ -51,7 +51,9 @@ void Copy(const Context& dev_ctx,
   VLOG(3) << "TensorCopy " << src.dims() << " from " << src.place() << " to "
           << dst_place;
 
-  dst->Resize(src.dims());
+  dst->set_meta(src.meta());
+  VLOG(3) << "src.numel() = " << src.numel()
+          << ", dst->numel() = " << dst->numel();
 
   void* dst_ptr = nullptr;
   if (paddle::platform::is_cpu_place(dst_place)) {

paddle/phi/kernels/cpu/add_n_kernel.cc

@@ -21,6 +21,13 @@ void AddNKernel(const Context& dev_ctx,
                 const std::vector<const TensorBase*>& x,
                 DenseTensor* out) {
   size_t in_num = x.size();
+  for (const TensorBase* tb : x) {
+    if (tb->initialized() && DenseTensor::classof(tb)) {
+      auto* dt = static_cast<const DenseTensor*>(tb);
+      out->set_meta(dt->meta());
+      break;
+    }
+  }
   dev_ctx.template Alloc<T>(out);
 
   bool in_place = false;

paddle/phi/kernels/gpu/add_n_kernel.cu

@@ -94,6 +94,15 @@ void AddNKernel(const Context &dev_ctx,
     grids = dim3(CEIL_DIV(length, tile_size), 1, 1);
     blocks = dim3(tile_size, 1, 1);
   };
+
+  for (const TensorBase *tb : x) {
+    if (tb->initialized() && DenseTensor::classof(tb)) {
+      auto *dt = static_cast<const DenseTensor *>(tb);
+      out->set_meta(dt->meta());
+      break;
+    }
+  }
+
   auto *out_ptr = dev_ctx.template Alloc<T>(out);
   bool in_place = false;
   if (x.size() > 0 && x[0]->initialized() && DenseTensor::classof(x[0])) {

python/paddle/fluid/backward.py

@@ -390,7 +390,7 @@ def _create_loss_op_desc_(loss):
         {},
         {"Out": [_append_grad_suffix_(loss.name)]},
         {
-            "shape": [1],
+            "shape": list(loss.shape),
             "value": 1.0,
             "dtype": loss.dtype,
             "force_cpu": False,

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

@@ -103,6 +103,7 @@ def false_func():
         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 @@ def false_func():
             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 @@ def test_0d_tensor_backward(self):
             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 @@ def greater_equal_branch(i, a):
         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()

python/paddle/static/nn/control_flow.py

@@ -969,7 +969,7 @@ def false_func():
     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):