[0-size Tensor No.362] Add 0-size Tensor support for where (#72934)

co63oc · web-flow · commit 993cf9b01097 · 2025-05-29T20:41:37.000+08:00
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -5673,7 +5673,8 @@ void WhereInferMeta(const MetaTensor& condition,
           x_dims.size()));
   size_t cond_dims_size = static_cast<size_t>(cond_dims.size());
   for (size_t i = 0; i < cond_dims_size; ++i) {
-    if (cond_dims[i] == -1 || x_dims[i] == -1) {
+    if (cond_dims[i] == -1 || x_dims[i] == -1 || cond_dims[i] == 0 ||
+        x_dims[i] == 0) {
       continue;
     }
     PADDLE_ENFORCE_EQ(
@@ -5696,7 +5697,8 @@ void WhereInferMeta(const MetaTensor& condition,
                         y_dims.size()));
   size_t x_dims_size = static_cast<size_t>(x_dims.size());
   for (size_t i = 0; i < x_dims_size; ++i) {
-    if (x_dims[i] == -1 || y_dims[i] == -1) {
+    if (x_dims[i] == -1 || y_dims[i] == -1 || x_dims[i] == 0 ||
+        y_dims[i] == 0) {
       continue;
     }
     PADDLE_ENFORCE_EQ(
diff --git a/paddle/phi/kernels/cpu/where_grad_kernel.cc b/paddle/phi/kernels/cpu/where_grad_kernel.cc
@@ -15,7 +15,7 @@
 #include "paddle/phi/kernels/where_grad_kernel.h"
 
 #include "paddle/phi/core/kernel_registry.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 
 template <typename T, typename Context>
@@ -29,6 +29,21 @@ void WhereGradKernel(const Context& ctx,
   const auto* cond_data = condition.data<bool>();
   auto numel = condition.numel();
   auto* dout = out_grad.data<T>();
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(ctx,
+                            phi::IntArray(common::vectorize(x_grad->dims())),
+                            static_cast<T>(0),
+                            x_grad);
+    }
+    if (y_grad) {
+      phi::Full<T, Context>(ctx,
+                            phi::IntArray(common::vectorize(y_grad->dims())),
+                            static_cast<T>(0),
+                            y_grad);
+    }
+    return;
+  }
 
   if (x_grad != nullptr) {
     auto* dx = ctx.template Alloc<T>(x_grad);
diff --git a/paddle/phi/kernels/cpu/where_kernel.cc b/paddle/phi/kernels/cpu/where_kernel.cc
@@ -30,6 +30,9 @@ void WhereKernel(const Context& ctx,
   auto x_numel = x.numel();
 
   T* out_data = ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) {
+    return;
+  }
 
   for (int i = 0; i < x_numel; i++) {
     out_data[i] = cond_data[i] ? x_data[i] : y_data[i];
diff --git a/paddle/phi/kernels/gpu/where_grad_kernel.cu b/paddle/phi/kernels/gpu/where_grad_kernel.cu
@@ -16,7 +16,7 @@
 
 #include "paddle/phi/backends/gpu/gpu_launch_config.h"
 #include "paddle/phi/core/kernel_registry.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 
 template <typename T, typename IndexT>
@@ -44,7 +44,21 @@ void WhereGradKernel(const Context& ctx,
   const bool* cond_data = condition.data<bool>();
   auto numel = condition.numel();
   auto* dout = out_grad.data<T>();
-
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(ctx,
+                            phi::IntArray(common::vectorize(x_grad->dims())),
+                            static_cast<T>(0),
+                            x_grad);
+    }
+    if (y_grad) {
+      phi::Full<T, Context>(ctx,
+                            phi::IntArray(common::vectorize(y_grad->dims())),
+                            static_cast<T>(0),
+                            y_grad);
+    }
+    return;
+  }
   T* dx = (x_grad != nullptr) ? ctx.template Alloc<T>(x_grad) : nullptr;
   T* dy = (y_grad != nullptr) ? ctx.template Alloc<T>(y_grad) : nullptr;
 
diff --git a/paddle/phi/kernels/gpu/where_kernel.cu b/paddle/phi/kernels/gpu/where_kernel.cu
@@ -38,6 +38,9 @@ void WhereKernel(const Context& ctx,
   std::vector<const DenseTensor*> ins = {&condition, &x, &y};
   std::vector<DenseTensor*> outs = {out};
   ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) {
+    return;
+  }
 
   CondFunctor<T> func;
   funcs::ElementwiseKernel<T, CondFunctor<T>, 1>(ctx, ins, &outs, func);
diff --git a/paddle/phi/kernels/xpu/where_grad_kernel.cc b/paddle/phi/kernels/xpu/where_grad_kernel.cc
@@ -16,7 +16,7 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 
 template <typename T, typename Context>
@@ -28,6 +28,22 @@ void WhereGradKernel(const Context& ctx,
                      DenseTensor* x_grad,
                      DenseTensor* y_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(ctx,
+                            phi::IntArray(common::vectorize(x_grad->dims())),
+                            static_cast<T>(0),
+                            x_grad);
+    }
+    if (y_grad) {
+      phi::Full<T, Context>(ctx,
+                            phi::IntArray(common::vectorize(y_grad->dims())),
+                            static_cast<T>(0),
+                            y_grad);
+    }
+    return;
+  }
+
   const auto* cond_data = condition.data<bool>();
   auto* dout = out_grad.data<T>();
 
diff --git a/paddle/phi/kernels/xpu/where_kernel.cc b/paddle/phi/kernels/xpu/where_kernel.cc
@@ -30,6 +30,9 @@ void WhereKernel(const Context& ctx,
   const XPUType* x_data = reinterpret_cast<const XPUType*>(x.data<T>());
   const XPUType* y_data = reinterpret_cast<const XPUType*>(y.data<T>());
   XPUType* out_data = reinterpret_cast<XPUType*>(ctx.template Alloc<T>(out));
+  if (out && out->numel() == 0) {
+    return;
+  }
 
   auto cond_dims = common::vectorize<int64_t>(condition.dims());
   auto x_dims = common::vectorize<int64_t>(x.dims());
diff --git a/test/legacy_test/test_where_op.py b/test/legacy_test/test_where_op.py
@@ -147,6 +147,13 @@ def init_config(self):
         self.cond = np.array(np.random.randint(2, size=(20, 2, 4)), dtype=bool)
 
 
+class TestWhereOp_ZeroSize(TestWhereOp):
+    def init_config(self):
+        self.x = np.random.uniform((-5), 5, (60, 0)).astype('float64')
+        self.y = np.random.uniform((-5), 5, (60, 0)).astype('float64')
+        self.cond = np.ones((60, 0)).astype('bool')
+
+
 class TestWhereAPI(unittest.TestCase):
     def setUp(self):
         self.init_data()
