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

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_kernel.cc

@@ -140,7 +140,11 @@ 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);
+    return;
+  }
   // axis < 0, cacluate the real axis
   if (axis < 0) {
     axis += in_dims.size();

paddle/phi/kernels/gpu/top_k_kernel.cu

@@ -61,6 +61,11 @@ 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);
+    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,13 @@ void ExpandAsGradKernel(const Context& context,
                         const std::vector<int>& target_shape,
                         DenseTensor* in_grad) {
   auto x_dims = x.dims();
+
+  if (x_dims.size() == 0) {
+    context.template Alloc<T>(in_grad);
+    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);

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;

paddle/phi/kernels/xpu/expand_as_kernel.cc

@@ -49,10 +49,27 @@ void ExpandAs(const Context& context,
               target_shape[i]));
     }
   }
+  if (target_shape.size() == 0) {
+    std::cout << "XYY Debug, debug the 「Tensor holds no memory」" << std::endl;
+    phi::DDim out_dims = phi::make_ddim(target_shape);
+    out->Resize(out_dims);
+    context.template Alloc<T>(out);
+
+    int r = xpu::copy<XPUType>(context.x_context(),
+                               reinterpret_cast<const XPUType*>(x.data<T>()),
+                               reinterpret_cast<XPUType*>(out->data<T>()),
+                               x.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
+    return;
+  }
+
   phi::DDim out_dims = phi::make_ddim(target_shape);
   out->Resize(out_dims);
   context.template Alloc<T>(out);
   auto& x_shape = vec_in_dims;
+  if (x.dims().size() == 0) {
+    x_shape = std::vector<int>({1});
+  }
   auto out_shape = phi::vectorize<int>(out_dims);
 
   int r = XPU_SUCCESS;
@@ -95,7 +112,7 @@ void ExpandAsKernel(const Context& ctx,
                         rank));
   PADDLE_ENFORCE_GE(
       rank,
-      1,
+      0,
       phi::errors::InvalidArgument("The rank (%d) of the input 'x' for "
                                    "expand_as_v2 op must be positive.",
                                    rank));

paddle/phi/kernels/xpu/expand_kernel.cc

@@ -78,7 +78,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_npu op must be positive, "
           "but the value received is %d.",
@@ -94,10 +94,21 @@ void ExpandKernel(const Context& ctx,
           shape_size,
           rank));
 
+  if (shape_size == 0) {
+    int r = xpu::copy<XPUType>(ctx.x_context(),
+                               reinterpret_cast<const XPUType*>(x.data<T>()),
+                               reinterpret_cast<XPUType*>(out->data<T>()),
+                               x.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
+    return;
+  }
   DDim out_dims = phi::make_ddim(final_expand_shape);
   out->Resize(out_dims);
   ctx.template Alloc<T>(out);
   auto& x_shape = vec_in_dims;
+  if (rank == 0) {
+    x_shape = std::vector<int>({1});
+  }
   auto out_shape = phi::vectorize<int>(out_dims);
 
   int r = XPU_SUCCESS;

paddle/phi/kernels/xpu/top_k_kernel.cc

@@ -49,7 +49,14 @@ void TopkKernel(const Context& dev_ctx,
       errors::External(
           "XPU API does not support smallest topk operation currently."
           " Operator will be supported in future update."));
-
+  if (in_dims.size() == 0) {
+    int r = xpu::copy<XPUType>(dev_ctx.x_context(),
+                               reinterpret_cast<const XPUType*>(x.data<T>()),
+                               reinterpret_cast<XPUType*>(out->data<T>()),
+                               x.numel());
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
+    return;
+  }
   if (axis < 0) axis += in_dims.size();
 
   size_t k = k_scalar.to<int>();