[0-size Tensor] Add 0-size Tensor support for paddle.multiply、paddle.__mul___ (#72984)

DanielSun11 · web-flow · commit cdb429983195 · 2025-05-30T16:09:10.000+08:00
* fix mul grad to support 0-size Tensor

* fix xpu kernel

* fix xpu error
diff --git a/paddle/phi/kernels/gpu/elementwise_grad.h b/paddle/phi/kernels/gpu/elementwise_grad.h
@@ -18,11 +18,11 @@ limitations under the License. */
 
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/broadcast_function.h"
 #include "paddle/phi/kernels/funcs/elementwise_grad_base.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
-
 namespace phi {
 
 template <typename T>
@@ -377,6 +377,26 @@ void ElementwiseMulGrad(const GPUContext &dev_ctx,
                         int axis) {
   const auto place = dev_ctx.GetPlace();
 
+  if (dout.numel() == 0) {
+    if (dx) {
+      if (dx->numel() == 0) {
+        dev_ctx.template Alloc<T>(dx);
+      } else {
+        phi::Full<T, GPUContext>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      if (dy->numel() == 0) {
+        dev_ctx.template Alloc<T>(dy);
+      } else {
+        phi::Full<T, GPUContext>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
+
   if (dx != nullptr && dy != nullptr) {
     std::vector<const DenseTensor *> ins = {&dout, &y, &x};
     GetGradXAndYOut<T>(dev_ctx,
diff --git a/paddle/phi/kernels/xpu/elementwise.h b/paddle/phi/kernels/xpu/elementwise.h
@@ -22,8 +22,8 @@ limitations under the License. */
 #include "paddle/phi/backends/xpu/xpu_context.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "xpu/refactor/math.h"
-
 namespace phi {
 
 template <typename T, typename XPUType>
@@ -112,6 +112,9 @@ void XPUElementwise(const XPUContext& dev_ctx,
                                       const std::vector<int64_t>&,
                                       const std::vector<int64_t>&)> func) {
   dev_ctx.template Alloc<T>(z);
+  if (z->numel() == 0) {
+    return;
+  }
   const DDim& x_dims = x.dims();
   const DDim& y_dims = y.dims();
 
@@ -194,6 +197,13 @@ void XPUElementwiseGrad(const XPUContext& dev_ctx,
   if (dy) {
     dy_data = dev_ctx.template Alloc<T>(dy);
   }
+  if (dz.numel() == 0) {
+    phi::Full<T, XPUContext>(
+        dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+    phi::Full<T, XPUContext>(
+        dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+    return;
+  }
 
   // use [1] to replace [], because xpu not support []
   if (x_dims_vec.size() == 0) {
diff --git a/test/legacy_test/test_multiply.py b/test/legacy_test/test_multiply.py
@@ -234,5 +234,61 @@ def multiply_shape_error():
         paddle.enable_static()
 
 
+class TestMultiplyApiZeroSize(TestMultiplyApi):
+
+    # only support the 0 size tensor
+    def _test_grad(self, x_data, y_data):
+        paddle.disable_static()
+        x = paddle.to_tensor(x_data, stop_gradient=False)
+        y = paddle.to_tensor(y_data, stop_gradient=False)
+        z = paddle.multiply(x, y)
+        loss = z.sum()
+        loss.backward()
+        np.testing.assert_allclose(
+            x.grad.numpy(), np.zeros(self.x_shape).astype('float32'), rtol=1e-05
+        )
+        np.testing.assert_allclose(
+            y.grad.numpy(), np.zeros(self.y_shape).astype('float32'), rtol=1e-05
+        )
+
+    def init_shapes(self):
+        self.x_shape = [0, 4]
+        self.y_shape = [0, 1]
+
+    def test_multiply(self):
+        np.random.seed(7)
+        self.init_shapes()
+
+        # test static computation graph
+        x_data = np.random.rand(*(self.x_shape)).astype('float32')
+        y_data = np.random.rand(*(self.y_shape)).astype('float32')
+        expected_res = np.multiply(x_data, y_data)
+        res = self._run_static_graph_case(x_data, y_data)
+        np.testing.assert_allclose(res, expected_res, rtol=1e-05)
+        # test dynamic computation graph
+        res = self._run_dynamic_graph_case(x_data, y_data)
+        np.testing.assert_allclose(res, expected_res, rtol=1e-05)
+        # test gradient
+        self._test_grad(x_data, y_data)
+
+
+class TestMultiplyApiZeroSize1(TestMultiplyApiZeroSize):
+    def init_shapes(self):
+        self.x_shape = [6, 0]
+        self.y_shape = [6, 0]
+
+
+class TestMultiplyApiZeroSize2(TestMultiplyApiZeroSize):
+    def init_shapes(self):
+        self.x_shape = [1, 8]
+        self.y_shape = [0, 1]
+
+
+class TestMultiplyApiZeroSize3(TestMultiplyApiZeroSize):
+    def init_shapes(self):
+        self.x_shape = [5, 0]
+        self.y_shape = [5, 1]
+
+
 if __name__ == '__main__':
     unittest.main()
