Fix (#73108)

Author: co63oc

paddle/phi/kernels/impl/bmm_grad_kernel_impl.h

@@ -54,6 +54,18 @@ void BmmGradKernel(const Context& dev_ctx,
                    const DenseTensor& out_grad,
                    DenseTensor* x_grad,
                    DenseTensor* y_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(y.dims())), 0, y_grad);
+    return;
+  }
+  if (y_grad && y_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(y_grad);
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(x.dims())), 0, x_grad);
+    return;
+  }
   DenseTensor x_help = x;
   DenseTensor y_help = y;
   DenseTensor out_grad_help = out_grad;

paddle/phi/kernels/xpu/bmm_grad_kernel.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/phi/kernels/bmm_grad_kernel.h"
 
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/xpu/bmm_xpu_utils.h"
 
 namespace phi {
@@ -60,6 +61,18 @@ void BmmGradKernel(const Context& dev_ctx,
                    const DenseTensor& out_grad,
                    DenseTensor* x_grad,
                    DenseTensor* y_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(y.dims())), 0, y_grad);
+    return;
+  }
+  if (y_grad && y_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(y_grad);
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(x.dims())), 0, x_grad);
+    return;
+  }
   DenseTensor x_help = x;
   DenseTensor y_help = y;
   DenseTensor out_grad_help = out_grad;

test/legacy_test/test_bmm_op.py

@@ -156,5 +156,23 @@ def test_api_error(self):
         self.assertRaises(ValueError, paddle.bmm, x_data, y_data_wrong3)
 
 
+class TestBmmOp_ZeroSize(OpTest):
+    def setUp(self):
+        self.op_type = "bmm"
+        self.python_api = paddle.bmm
+        self.public_python_api = paddle.bmm
+        X = np.random.random((10, 0, 4)).astype("float64")
+        Y = np.random.random((10, 4, 5)).astype("float64")
+        self.inputs = {'X': X, 'Y': Y}
+        Out = np.matmul(X, Y)
+        self.outputs = {'Out': Out}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+    def test_checkout_grad(self):
+        self.check_grad(['X', 'Y'], 'Out', check_pir=True)
+
+
 if __name__ == "__main__":
     unittest.main()