From 78c7c420d0b8334c027b12a5f0e593bcdee1e6be Mon Sep 17 00:00:00 2001
From: co63oc <co63oc@users.noreply.github.com>
Date: Thu, 5 Jun 2025 11:11:57 +0800
Subject: [PATCH] Fix

---
 paddle/phi/kernels/impl/bmm_grad_kernel_impl.h | 12 ++++++++++++
 paddle/phi/kernels/xpu/bmm_grad_kernel.cc      | 13 +++++++++++++
 test/legacy_test/test_bmm_op.py                | 18 ++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/paddle/phi/kernels/impl/bmm_grad_kernel_impl.h b/paddle/phi/kernels/impl/bmm_grad_kernel_impl.h
index 33ca7233a37268..b05caf497cd601 100644
--- a/paddle/phi/kernels/impl/bmm_grad_kernel_impl.h
+++ b/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;
diff --git a/paddle/phi/kernels/xpu/bmm_grad_kernel.cc b/paddle/phi/kernels/xpu/bmm_grad_kernel.cc
index e2fdbb610d2a2b..f533ac0f97b82c 100644
--- a/paddle/phi/kernels/xpu/bmm_grad_kernel.cc
+++ b/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;
diff --git a/test/legacy_test/test_bmm_op.py b/test/legacy_test/test_bmm_op.py
index 132a227c0ac443..a8ee7955f6375b 100644
--- a/test/legacy_test/test_bmm_op.py
+++ b/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()