[Zero-Dim] Support 0D Tensor input for topk/broadcast_to/expand/expand_as/broadcast_shape (#50536)

paddle/phi/infermeta/unary.cc

@@ -472,6 +472,7 @@ void CumInferMeta(const MetaTensor& x,
     out->set_dims(x_dims);
     out->set_dtype(x.dtype());
   }
+
   out->share_lod(x);
 }
 
@@ -970,7 +971,7 @@ void ExpandInferMeta(const MetaTensor& x,
                                    MAX_RANK_SUPPORTED));
   PADDLE_ENFORCE_GE(
       expand_shape.size(),
-      1,
+      0,
       phi::errors::InvalidArgument("The number of elements (%d) of 'shape' for "
                                    "must be a positive integer.",
                                    expand_shape.size()));
@@ -1005,7 +1006,7 @@ void ExpandInferMeta(const MetaTensor& x,
 
   out->set_dims(make_ddim(out_shape));
   out->set_dtype(x.dtype());
-  if (out_shape[0] == x_dims[0]) {
+  if (out_rank > 0 && out_shape[0] == x_dims[0]) {
     out->share_lod(x);
   }
 }
@@ -4097,14 +4098,23 @@ void TopKInferMeta(const MetaTensor& x,
                    MetaConfig config) {
   auto input_dims = x.dims();
   const int& dim_size = input_dims.size();
-  PADDLE_ENFORCE_EQ(
-      (axis < dim_size) && (axis >= (-1 * dim_size)),
-      true,
-      phi::errors::InvalidArgument(
-          "the axis of topk must be [-%d, %d), but you set axis is %d",
-          dim_size,
-          dim_size,
-          axis));
+  if (dim_size != 0) {
+    PADDLE_ENFORCE_EQ(
+        (axis < dim_size) && (axis >= (-1 * dim_size)),
+        true,
+        phi::errors::InvalidArgument(
+            "the axis of topk must be [-%d, %d), but you set axis is %d",
+            dim_size,
+            dim_size,
+            axis));
+  } else {
+    PADDLE_ENFORCE_EQ(
+        (axis == dim_size) || (axis == -1),
+        true,
+        phi::errors::InvalidArgument("the axis of topk must be 0 or -1 when "
+                                     "x.dims() = 0, but you set axis is %d",
+                                     axis));
+  }
 
   if (axis < 0) axis += dim_size;
 
@@ -4122,12 +4132,13 @@ void TopKInferMeta(const MetaTensor& x,
 
   PADDLE_ENFORCE_GE(
       input_dims.size(),
-      1,
-      phi::errors::InvalidArgument("input of topk must have >= 1d shape"));
+      0,
+      phi::errors::InvalidArgument("input of topk must have >= 0d shape"));
 
   phi::DDim dims = input_dims;
-
-  dims[axis] = k;
+  if (input_dims.size() > 0) {
+    dims[axis] = k;
+  }
   out->set_dims(dims);
   out->share_lod(x);
   out->set_dtype(x.dtype());

paddle/phi/kernels/cpu/top_k_grad_kernel.cc

@@ -66,6 +66,11 @@ void TopkGradKernel(const Context& dev_ctx,
   axis = (axis < 0) ? (in_dims.size() + axis) : axis;
 
   T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
+  if (in_dims.size() == 0) {
+    phi::Copy<Context>(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    return;
+  }
+
   if (axis + 1 == in_dims.size()) {
     // allocate the memory for the input_grad
 

paddle/phi/kernels/cpu/top_k_kernel.cc

@@ -140,7 +140,13 @@ void TopkKernel(const Context& dev_ctx,
   const auto* input = &x;
   // Get the top k elements of each row of input tensor
   const auto& in_dims = input->dims();
-
+  // 0d input x
+  if (in_dims.size() == 0) {
+    phi::Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+    dev_ctx.template Alloc<int64_t>(indices);
+    phi::funcs::set_constant(dev_ctx, indices, 0.0);
+    return;
+  }
   // axis < 0, cacluate the real axis
   if (axis < 0) {
     axis += in_dims.size();

paddle/phi/kernels/gpu/top_k_grad_kernel.cu

@@ -47,6 +47,11 @@ void TopkGradKernel(const Context& dev_ctx,
   const T* out_grad_data = out_grad.data<T>();
   const int64_t* indices_data = indices.data<int64_t>();
 
+  if (in_dims.size() == 0) {
+    phi::Copy<Context>(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    return;
+  }
+
   int pre, n, post;
   phi::funcs::GetDims(in_dims, axis, &pre, &n, &post);
 

paddle/phi/kernels/gpu/top_k_kernel.cu

@@ -61,6 +61,14 @@ void TopkKernel(const Context& dev_ctx,
   const auto* input = &x;
   // get the input dims
   const auto& in_dims = input->dims();
+
+  // 0d input tensor
+  if (in_dims.size() == 0) {
+    phi::Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+    dev_ctx.template Alloc<int64_t>(indices);
+    phi::funcs::set_constant(dev_ctx, indices, 0.0);
+    return;
+  }
   // calcluate the real axis
   if (axis < 0) axis += in_dims.size();
 

paddle/phi/kernels/impl/expand_as_grad_kernel_impl.h

@@ -49,6 +49,12 @@ void ExpandAsGradKernel(const Context& context,
                         const std::vector<int>& target_shape,
                         DenseTensor* in_grad) {
   auto x_dims = x.dims();
+
+  if (in_grad->dims() == out_grad.dims()) {
+    phi::Copy(context, out_grad, context.GetPlace(), false, in_grad);
+    return;
+  }
+
   auto vec_in_dims = phi::vectorize<int>(x_dims);
   auto diff = target_shape.size() - vec_in_dims.size();
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
@@ -65,64 +71,56 @@ void ExpandAsGradKernel(const Context& context,
   }
 
   int dims = reduce_dims_vec.size();
-  bool just_copy = true;
-  for (size_t i = 0; i < repeat_times.size(); i++) {
-    if (repeat_times[i] != 1) {
-      just_copy = false;
+
+  PADDLE_ENFORCE_GE(
+      dims,
+      0,
+      errors::InvalidArgument("The rank of the input 'Out@GRAD' for "
+                              "expand_as_v2_grad op must be greater than or "
+                              "equal to 0, but the value received is %d.",
+                              dims));
+  PADDLE_ENFORCE_LE(
+      dims,
+      MAX_RANK_SUPPORTED,
+      errors::InvalidArgument("The rank of the input 'Out@GRAD' for "
+                              "expand_as_v2_grad op must be less than or equal "
+                              "to %d, but the value received is %d.",
+                              MAX_RANK_SUPPORTED,
+                              dims));
+  switch (dims) {
+    case 0:
+      ExpandAsBackward<Context, T, 0>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
       break;
-    }
-  }
-  // no need reduce, just copy
-  if (just_copy) {
-    context.template Alloc<T>(in_grad);
-    phi::Copy(context, out_grad, context.GetPlace(), false, in_grad);
-  } else {
-    PADDLE_ENFORCE_GE(
-        dims,
-        1,
-        errors::InvalidArgument("The rank of the input 'Out@GRAD' for "
-                                "expand_as_v2_grad op must be greater than or "
-                                "equal to 1, but the value received is %d.",
-                                dims));
-    PADDLE_ENFORCE_LE(dims,
-                      MAX_RANK_SUPPORTED,
-                      errors::InvalidArgument(
-                          "The rank of the input 'Out@GRAD' for "
-                          "expand_as_v2_grad op must be less than or equal "
-                          "to %d, but the value received is %d.",
-                          MAX_RANK_SUPPORTED,
-                          dims));
-    switch (dims) {
-      case 1:
-        ExpandAsBackward<Context, T, 1>(
-            context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 2:
-        ExpandAsBackward<Context, T, 2>(
-            context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 3:
-        ExpandAsBackward<Context, T, 3>(
-            context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 4:
-        ExpandAsBackward<Context, T, 4>(
-            context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 5:
-        ExpandAsBackward<Context, T, 5>(
-            context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 6:
-        ExpandAsBackward<Context, T, 6>(
-            context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      default:
-        PADDLE_THROW(errors::InvalidArgument(
-            "Only support tensor with rank being between 1 and 6. But "
-            "received tensor's rank = %d.",
-            dims));
-    }
+    case 1:
+      ExpandAsBackward<Context, T, 1>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 2:
+      ExpandAsBackward<Context, T, 2>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 3:
+      ExpandAsBackward<Context, T, 3>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 4:
+      ExpandAsBackward<Context, T, 4>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 5:
+      ExpandAsBackward<Context, T, 5>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 6:
+      ExpandAsBackward<Context, T, 6>(
+          context, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    default:
+      PADDLE_THROW(errors::InvalidArgument(
+          "Only support tensor with rank being between 1 and 6. But "
+          "received tensor's rank = %d.",
+          dims));
   }
 }
 

paddle/phi/kernels/impl/expand_as_kernel_impl.h

@@ -34,6 +34,10 @@ void ExpandAs(const Context& context,
   auto diff = target_shape.size() - vec_in_dims.size();
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
   std::vector<int> repeat_times(vec_in_dims.size());
+  if (Rank == 0) {
+    phi::Copy<Context>(context, x, context.GetPlace(), false, out);
+    return;
+  }
   for (size_t i = 0; i < vec_in_dims.size(); ++i) {
     PADDLE_ENFORCE_NE(
         target_shape[i],
@@ -108,7 +112,7 @@ void ExpandAsKernel(const Context& ctx,
                         rank));
   PADDLE_ENFORCE_GE(
       rank,
-      1,
+      0,
       errors::InvalidArgument("The rank (%d) of the input 'x' for "
                               "expand_as_v2 op must be positive.",
                               rank));
@@ -133,6 +137,9 @@ void ExpandAsKernel(const Context& ctx,
   }
 
   switch (target_rank) {
+    case 0:
+      ExpandAs<Context, T, 0>(ctx, x, real_target_shape, out);
+      break;
     case 1:
       ExpandAs<Context, T, 1>(ctx, x, real_target_shape, out);
       break;

paddle/phi/kernels/impl/expand_grad_kernel_impl.h

@@ -54,6 +54,11 @@ void ExpandGradKernel(const Context& ctx,
                       DenseTensor* in_grad) {
   auto expand_shape = shape.GetData();
   auto x_dims = x.dims();
+
+  if (in_grad->dims() == out_grad.dims()) {
+    phi::Copy(ctx, out_grad, ctx.GetPlace(), false, in_grad);
+    return;
+  }
   auto vec_in_dims = phi::vectorize<int>(x_dims);
   auto diff = expand_shape.size() - vec_in_dims.size();
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
@@ -79,63 +84,55 @@ void ExpandGradKernel(const Context& ctx,
 
   int dims = reduce_dims_vec.size();
 
-  bool just_copy = true;
-  for (size_t i = 0; i < repeat_times.size(); i++) {
-    if (repeat_times[i] != 1) {
-      just_copy = false;
+  PADDLE_ENFORCE_GE(
+      dims,
+      0,
+      phi::errors::InvalidArgument("The rank of the input 'Out@GRAD' for "
+                                   "expand_v2_grad op must be greater than or "
+                                   "equal to 0, but the value received is %d.",
+                                   dims));
+  PADDLE_ENFORCE_LE(dims,
+                    MAX_RANK_SUPPORTED,
+                    phi::errors::InvalidArgument(
+                        "The rank of the input 'Out@GRAD' for "
+                        "expand_v2_grad op must be less than or equal "
+                        "to %d, but the value received is %d.",
+                        MAX_RANK_SUPPORTED,
+                        dims));
+  switch (dims) {
+    case 0:
+      ExpandBackward<Context, T, 1>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
       break;
-    }
-  }
-  // no need reduce, just copy
-  if (just_copy) {
-    phi::Copy(ctx, out_grad, ctx.GetPlace(), false, in_grad);
-  } else {
-    PADDLE_ENFORCE_GE(dims,
-                      1,
-                      phi::errors::InvalidArgument(
-                          "The rank of the input 'Out@GRAD' for "
-                          "expand_v2_grad op must be greater than or "
-                          "equal to 1, but the value received is %d.",
-                          dims));
-    PADDLE_ENFORCE_LE(dims,
-                      MAX_RANK_SUPPORTED,
-                      phi::errors::InvalidArgument(
-                          "The rank of the input 'Out@GRAD' for "
-                          "expand_v2_grad op must be less than or equal "
-                          "to %d, but the value received is %d.",
-                          MAX_RANK_SUPPORTED,
-                          dims));
-    switch (dims) {
-      case 1:
-        ExpandBackward<Context, T, 1>(
-            ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 2:
-        ExpandBackward<Context, T, 2>(
-            ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 3:
-        ExpandBackward<Context, T, 3>(
-            ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 4:
-        ExpandBackward<Context, T, 4>(
-            ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 5:
-        ExpandBackward<Context, T, 5>(
-            ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      case 6:
-        ExpandBackward<Context, T, 6>(
-            ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
-        break;
-      default:
-        PADDLE_THROW(phi::errors::InvalidArgument(
-            "Only support tensor with rank being between 1 and 6. But "
-            "received tensor's rank = %d.",
-            dims));
-    }
+    case 1:
+      ExpandBackward<Context, T, 1>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 2:
+      ExpandBackward<Context, T, 2>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 3:
+      ExpandBackward<Context, T, 3>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 4:
+      ExpandBackward<Context, T, 4>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 5:
+      ExpandBackward<Context, T, 5>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    case 6:
+      ExpandBackward<Context, T, 6>(
+          ctx, out_grad, reshape_dims_vec, reduce_dims_vec, in_grad);
+      break;
+    default:
+      PADDLE_THROW(phi::errors::InvalidArgument(
+          "Only support tensor with rank being between 1 and 6. But "
+          "received tensor's rank = %d.",
+          dims));
   }
 }
 

paddle/phi/kernels/impl/expand_kernel_impl.h

@@ -35,6 +35,10 @@ void Expand(const Context& ctx,
   auto diff = expand_shape.size() - vec_in_dims.size();
   vec_in_dims.insert(vec_in_dims.begin(), diff, 1);
   std::vector<int> repeat_times(vec_in_dims.size());
+  if (Rank == 0) {
+    phi::Copy<Context>(ctx, x, ctx.GetPlace(), false, out);
+    return;
+  }
   for (size_t i = 0; i < vec_in_dims.size(); ++i) {
     PADDLE_ENFORCE_NE(
         expand_shape[i],
@@ -74,7 +78,6 @@ void Expand(const Context& ctx,
       repeat_times[i] = 1;
     }
   }
-
   Eigen::DSizes<Eigen::DenseIndex, Rank> bcast_dims;
   for (size_t i = 0; i < repeat_times.size(); ++i) {
     bcast_dims[i] = repeat_times[i];
@@ -112,7 +115,7 @@ void ExpandKernel(const Context& ctx,
   auto rank = x.dims().size();
   PADDLE_ENFORCE_GE(
       rank,
-      1,
+      0,
       phi::errors::InvalidArgument(
           "The rank of the input 'X' for expand_v2 op must be positive, "
           "but the value received is %d.",
@@ -145,6 +148,9 @@ void ExpandKernel(const Context& ctx,
           MAX_RANK_SUPPORTED));
   rank = std::max(rank, static_cast<int>(shape_size));
   switch (rank) {
+    case 0:
+      Expand<Context, T, 0>(ctx, x, shape, out);
+      break;
     case 1:
       Expand<Context, T, 1>(ctx, x, shape, out);
       break;

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

@@ -559,6 +559,210 @@ class TestSundryAPI(unittest.TestCase):
         paddle.disable_static()
         self.x = paddle.rand([])
 
+    def test_expand(self):
+        # case1
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        out = paddle.expand(x, shape=[1])
+        out.retain_grads()
+        out.backward()
+
+        self.assertEqual(out.shape, [1])
+        np.testing.assert_allclose(out, 1.0)
+        self.assertEqual(x.grad.shape, [])
+        np.testing.assert_allclose(x.grad, 1.0)
+        self.assertEqual(out.grad.shape, [1])
+        np.testing.assert_allclose(out.grad, 1.0)
+
+        # case2
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        out1 = paddle.expand(x1, shape=[])
+        out1.retain_grads()
+        out1.backward()
+
+        self.assertEqual(out1.shape, [])
+        np.testing.assert_allclose(out1, 1.0)
+        self.assertEqual(x1.grad.shape, [])
+        np.testing.assert_allclose(x1.grad, 1.0)
+        self.assertEqual(out1.grad.shape, [])
+        np.testing.assert_allclose(out1.grad, 1.0)
+
+        # case3
+        x2 = paddle.full([], 1, 'float32')
+        x2.stop_gradient = False
+        out2 = paddle.expand(x2, shape=[1, 1])
+        out2.retain_grads()
+        out2.backward()
+
+        self.assertEqual(out2.shape, [1, 1])
+        np.testing.assert_allclose(out2, 1.0)
+        self.assertEqual(x2.grad.shape, [])
+        np.testing.assert_allclose(x2.grad, 1.0)
+        self.assertEqual(out2.grad.shape, [1, 1])
+        np.testing.assert_allclose(out2.grad, 1.0)
+
+        # case4
+        x3 = paddle.full([], 1, 'float32')
+        x3.stop_gradient = False
+        out3 = paddle.expand(x3, shape=[3, 3])
+        out3.retain_grads()
+        out3.backward()
+
+        self.assertEqual(out3.shape, [3, 3])
+        np.testing.assert_allclose(out3, 1.0)
+        self.assertEqual(x3.grad.shape, [])
+        np.testing.assert_allclose(x3.grad, 9.0)
+        self.assertEqual(out3.grad.shape, [3, 3])
+        np.testing.assert_allclose(out3.grad, 1.0)
+
+    def test_expand_as(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        y = paddle.full([], 1, 'float32')
+        y.stop_gradient = False
+        out = paddle.expand_as(x, y)
+        out.backward()
+        self.assertEqual(x.shape, [])
+        self.assertEqual(x.item(), 1.0)
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad.item(), 1.0)
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out.item(), 1.0)
+        self.assertEqual(out.grad, None)
+
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        y1 = paddle.full([1], 1, 'float32')
+        out1 = paddle.expand_as(x1, y1)
+        out1.backward()
+        self.assertEqual(x1.shape, [])
+        self.assertEqual(x1.item(), 1.0)
+        self.assertEqual(x1.grad.shape, [])
+        self.assertEqual(x1.grad.item(0), 1.0)
+        self.assertEqual(out1.shape, [1])
+        self.assertEqual(out1.item(0), 1.0)
+        self.assertEqual(out1.grad, None)
+
+        x2 = paddle.full([], 1, 'float32')
+        x2.stop_gradient = False
+        y2 = paddle.full([3, 3], 1, 'float32')
+        out2 = paddle.expand_as(x2, y2)
+        out2.backward()
+        self.assertEqual(x2.shape, [])
+        self.assertEqual(x2.item(), 1.0)
+        self.assertEqual(x2.grad.shape, [])
+        self.assertEqual(x2.grad.item(0), 9.0)
+        self.assertEqual(out2.shape, [3, 3])
+        self.assertEqual(out2.item(0), 1.0)
+        self.assertEqual(out2.grad, None)
+
+    def test_top_k(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        out, indices = paddle.topk(x, k=1, axis=0)
+        out.retain_grads()
+        out.backward()
+        self.assertEqual(indices.shape, [])
+        self.assertEqual(indices.item(), 0)
+        self.assertEqual(x.shape, [])
+        self.assertEqual(x.item(), 1.0)
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad.item(0), 1.0)
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out.item(), 1.0)
+        self.assertEqual(out.grad, 1.0)
+
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        out1, indices1 = paddle.topk(x1, k=1, axis=-1)
+        out1.retain_grads()
+        out1.backward()
+        self.assertEqual(indices1.shape, [])
+        self.assertEqual(indices1.item(), 0)
+        self.assertEqual(x1.shape, [])
+        self.assertEqual(x1.item(), 1.0)
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad.item(0), 1.0)
+        self.assertEqual(out1.shape, [])
+        self.assertEqual(out1.item(), 1.0)
+        self.assertEqual(out1.grad, 1.0)
+
+        with self.assertRaises(ValueError):
+            tmp = paddle.topk(x1, k=1, axis=2)
+
+    def test_broadcast_to(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        out = paddle.broadcast_to(x, shape=[1])
+        out.retain_grads()
+        out.backward()
+
+        self.assertEqual(out.shape, [1])
+        np.testing.assert_allclose(out, 1.0)
+        self.assertEqual(x.grad.shape, [])
+        np.testing.assert_allclose(x.grad, 1.0)
+        self.assertEqual(out.grad.shape, [1])
+        np.testing.assert_allclose(out.grad, 1.0)
+
+        # case2
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        out1 = paddle.broadcast_to(x1, shape=[])
+        out1.retain_grads()
+        out1.backward()
+
+        self.assertEqual(out1.shape, [])
+        np.testing.assert_allclose(out1, 1.0)
+        self.assertEqual(x1.grad.shape, [])
+        np.testing.assert_allclose(x1.grad, 1.0)
+        self.assertEqual(out1.grad.shape, [])
+        np.testing.assert_allclose(out1.grad, 1.0)
+
+        # case3
+        x2 = paddle.full([], 1, 'float32')
+        x2.stop_gradient = False
+        out2 = paddle.broadcast_to(x2, shape=[1, 1])
+        out2.retain_grads()
+        out2.backward()
+
+        self.assertEqual(out2.shape, [1, 1])
+        np.testing.assert_allclose(out2, 1.0)
+        self.assertEqual(x2.grad.shape, [])
+        np.testing.assert_allclose(x2.grad, 1.0)
+        self.assertEqual(out2.grad.shape, [1, 1])
+        np.testing.assert_allclose(out2.grad, 1.0)
+
+        # case4
+        x3 = paddle.full([], 1, 'float32')
+        x3.stop_gradient = False
+        out3 = paddle.broadcast_to(x3, shape=[3, 3])
+        out3.retain_grads()
+        out3.backward()
+
+        self.assertEqual(out3.shape, [3, 3])
+        np.testing.assert_allclose(out3, 1.0)
+        self.assertEqual(x3.grad.shape, [])
+        np.testing.assert_allclose(x3.grad, 9.0)
+        self.assertEqual(out3.grad.shape, [3, 3])
+        np.testing.assert_allclose(out3.grad, 1.0)
+
+    def test_broadcast_shape(self):
+        x = []
+        y = [3, 5]
+        out = paddle.broadcast_shape(x, y)
+        self.assertEqual(out, [3, 5])
+
+        x = [3, 5]
+        y = []
+        out = paddle.broadcast_shape(x, y)
+        self.assertEqual(out, [3, 5])
+
+        x = []
+        y = []
+        out = paddle.broadcast_shape(x, y)
+        self.assertEqual(out, [])
+
     def test_argmin(self):
         x = paddle.rand([])
         out1 = paddle.argmin(x, 0)
@@ -1673,6 +1877,200 @@ class TestSundryAPIStatic(unittest.TestCase):
         self.exe = paddle.static.Executor()
 
     @prog_scope()
+    def test_expand(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        out = paddle.expand(x, shape=[1])
+        paddle.static.append_backward(out.sum())
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x, out, x.grad_name, out.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, (1,))
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+        self.assertEqual(res[3].shape, (1,))
+        self.assertEqual(res[3], 1.0)
+
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        out1 = paddle.expand(x1, shape=[])
+        paddle.static.append_backward(out1.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x1, out1, x1.grad_name, out1.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+        self.assertEqual(res[3].shape, ())
+        self.assertEqual(res[3], 1.0)
+
+        x2 = paddle.full([], 1, 'float32')
+        x2.stop_gradient = False
+        out2 = paddle.expand(x2, shape=[3, 3])
+        paddle.static.append_backward(out2.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x2, out2, x2.grad_name, out2.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, (3, 3))
+        self.assertEqual(res[1].any(), 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 9)
+        self.assertEqual(res[3].shape, (3, 3))
+        self.assertEqual(res[3].any(), 1.0)
+
+    @prog_scope()
+    def test_expand_as(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        y = paddle.full([], 1, 'float32')
+        y.stop_gradient = False
+        out = paddle.expand_as(x, y)
+        paddle.static.append_backward(out.sum())
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x, out, x.grad_name, out.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+        self.assertEqual(res[3].shape, ())
+        self.assertEqual(res[3], 1.0)
+
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        y1 = paddle.full([1], 1, 'float32')
+        y1.stop_gradient = False
+        out1 = paddle.expand_as(x1, y1)
+        paddle.static.append_backward(out1.sum())
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x1, out1, x1.grad_name, out1.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, (1,))
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+        self.assertEqual(res[3].shape, (1,))
+        self.assertEqual(res[3], 1.0)
+
+        x2 = paddle.full([], 1, 'float32')
+        x2.stop_gradient = False
+        y2 = paddle.full([3, 3], 1, 'float32')
+        y2.stop_gradient = False
+        out2 = paddle.expand_as(x2, y2)
+        paddle.static.append_backward(out2.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x2, out2, x2.grad_name, out2.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, (3, 3))
+        self.assertEqual(res[1].any(), 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 9)
+        self.assertEqual(res[3].shape, (3, 3))
+        self.assertEqual(res[3].any(), 1.0)
+
+    @prog_scope()
+    def test_top_k(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        out, indices = paddle.topk(x, k=1, axis=0)
+        paddle.static.append_backward(out.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x, out, indices, x.grad_name, out.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 0.0)
+        self.assertEqual(res[3].shape, ())
+        self.assertEqual(res[3], 1.0)
+        self.assertEqual(res[4].shape, ())
+        self.assertEqual(res[4], 1.0)
+
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        out1, indices1 = paddle.topk(x1, k=1, axis=-1)
+        paddle.static.append_backward(out1.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x1, out1, indices1, x1.grad_name, out1.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 0.0)
+        self.assertEqual(res[3].shape, ())
+        self.assertEqual(res[3], 1.0)
+        self.assertEqual(res[4].shape, ())
+        self.assertEqual(res[4], 1.0)
+
+        with self.assertRaises(ValueError):
+            tmp = paddle.topk(x1, k=1, axis=2)
+
+    @prog_scope()
+    def test_broadcast_to(self):
+        x = paddle.full([], 1, 'float32')
+        x.stop_gradient = False
+        out = paddle.broadcast_to(x, shape=[1])
+        paddle.static.append_backward(out.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x, out, x.grad_name, out.grad_name]
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, (1,))
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+        self.assertEqual(res[3].shape, (1,))
+        self.assertEqual(res[3], 1.0)
+
+        x1 = paddle.full([], 1, 'float32')
+        x1.stop_gradient = False
+        out1 = paddle.broadcast_to(x1, shape=[])
+        paddle.static.append_backward(out1.sum())
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog, fetch_list=[x1, out1, x1.grad_name, out1.grad_name]
+        )
+
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[0], 1.0)
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+        self.assertEqual(res[3].shape, ())
+        self.assertEqual(res[3], 1.0)
+
     def test_argmin(self):
         x = paddle.rand([])
         out1 = paddle.argmin(x, 0)

python/paddle/tensor/manipulation.py

@@ -3192,7 +3192,7 @@ def expand_as(x, y, name=None):
 
     Expand the input tensor ``x`` to the same shape as the input tensor ``y``.
 
-    Both the number of dimensions of ``x`` and ``y`` must be less than or equal to 6, and the number of dimensions of ``y`` must be greather than or equal to that of ``x``. The dimension to expand must have a value of 1.
+    Both the number of dimensions of ``x`` and ``y`` must be less than or equal to 6, and the number of dimensions of ``y`` must be greather than or equal to that of ``x``. The dimension to expand must have a value of 0.
 
     Args:
         x (Tensor): The input tensor, its data type is bool, float32, float64, int32 or int64.
@@ -3252,13 +3252,13 @@ def broadcast_to(x, shape, name=None):
 
     Broadcast the input tensor to a given shape.
 
-    Both the number of dimensions of ``x`` and the number of elements in ``shape`` should be less than or equal to 6. The dimension to broadcast to must have a value 1.
+    Both the number of dimensions of ``x`` and the number of elements in ``shape`` should be less than or equal to 6. The dimension to broadcast to must have a value 0.
 
 
     Args:
         x (Tensor): The input tensor, its data type is bool, float32, float64, int32 or int64.
         shape (list|tuple|Tensor): The result shape after broadcasting. The data type is int32. If shape is a list or tuple, all its elements
-            should be integers or 1-D Tensors with the data type int32. If shape is a Tensor, it should be an 1-D Tensor with the data type int32.
+            should be integers or 0-D or 1-D Tensors with the data type int32. If shape is a Tensor, it should be an 1-D Tensor with the data type int32.
             The value -1 in shape means keeping the corresponding dimension unchanged.
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
     Returns:
@@ -3280,13 +3280,13 @@ def broadcast_to(x, shape, name=None):
         if isinstance(shape, Variable):
             assert len(shape.shape) == 1, 'shape must be an 1-D Tensor.'
         else:
+            type_tuple = (int, np.int32, np.int64)
             for elem in shape:
                 if isinstance(elem, Variable):
                     assert (
                         len(elem.shape) == 1
                     ), 'Elements in shape must be 1-D Tensors or integers.'
                 else:
-                    type_tuple = (int, np.int32, np.int64)
                     assert isinstance(
                         elem, type_tuple
                     ), 'Elements in shape must be 1-D Tensors or integers.'
@@ -3346,12 +3346,12 @@ def expand(x, shape, name=None):
 
     Expand the input tensor to a given shape.
 
-    Both the number of dimensions of ``x`` and the number of elements in ``shape`` should be less than or equal to 6. And the number of dimensions of ``x`` should be less than the number of elements in ``shape``. The dimension to expand must have a value 1.
+    Both the number of dimensions of ``x`` and the number of elements in ``shape`` should be less than or equal to 6. And the number of dimensions of ``x`` should be less than the number of elements in ``shape``. The dimension to expand must have a value 0.
 
     Args:
         x (Tensor): The input Tensor, its data type is bool, float32, float64, int32 or int64.
         shape (list|tuple|Tensor): The result shape after expanding. The data type is int32. If shape is a list or tuple, all its elements
-            should be integers or 1-D Tensors with the data type int32. If shape is a Tensor, it should be an 1-D Tensor with the data type int32.
+            should be integers or 0-D or 1-D Tensors with the data type int32. If shape is a Tensor, it should be an 1-D Tensor with the data type int32.
             The value -1 in shape means keeping the corresponding dimension unchanged.
         name (str, optional): The default value is None. Normally there is no need for user to set this property. For more information, please refer to :ref:`api_guide_Name` .