fix dist_grad kernel (#53239)

paddle/phi/kernels/dist_grad_kernel.cc

@@ -52,6 +52,10 @@ void DistGradKernel(const Context& dev_ctx,
                     float p,
                     DenseTensor* x_grad,
                     DenseTensor* y_grad) {
+  if ((!x_grad) && (!y_grad)) {
+    return;
+  }
+
   auto t = Subtract<T, Context>(dev_ctx, x, y);
   DenseTensor x_grad_tmp;
   x_grad_tmp.Resize(t.dims());
@@ -59,7 +63,8 @@ void DistGradKernel(const Context& dev_ctx,
   y_grad_tmp.Resize(t.dims());
   PNormGradKernel<T, Context>(
       dev_ctx, t, out, out_grad, p, -1, 1e-12, false, true, &x_grad_tmp);
-  ScaleKernel<T, Context>(dev_ctx, x_grad_tmp, -1.0, 0.0, false, &y_grad_tmp);
+
+  if (x_grad) {
     // do reduce, the implemetation of cpu SumKernel has bug, it changes
     // the dims of output iternally, so we Resize x/y_grad twice.
     auto res_x = GetReduceDims(x_grad_tmp.dims(), x.dims());
@@ -71,6 +76,10 @@ void DistGradKernel(const Context& dev_ctx,
     } else {
       x_grad->ShareBufferWith(x_grad_tmp);
     }
+  }
+
+  if (y_grad) {
+    ScaleKernel<T, Context>(dev_ctx, x_grad_tmp, -1.0, 0.0, false, &y_grad_tmp);
     auto res_y = GetReduceDims(y_grad_tmp.dims(), y.dims());
     if (!std::get<0>(res_y).empty()) {
       y_grad->Resize(phi::make_ddim(std::get<1>(res_y)));
@@ -80,6 +89,7 @@ void DistGradKernel(const Context& dev_ctx,
     } else {
       y_grad->ShareBufferWith(y_grad_tmp);
     }
+  }
 }
 
 }  // namespace phi

python/paddle/fluid/tests/unittests/test_dist_op.py

@@ -192,6 +192,15 @@ class TestDistAPI(unittest.TestCase):
             )
             np.testing.assert_allclose(dist(x_i, y_i, p), out[0], rtol=1e-05)
 
+    def test_grad_x(self):
+        paddle.disable_static()
+        a = paddle.rand([2, 2, 3, 2])
+        b = paddle.rand([1, 1, 3, 1])
+        a.stop_gradient = False
+        c = paddle.dist(a, b, 2)
+        c.backward()
+        paddle.enable_static()
+
 
 if __name__ == '__main__':
     paddle.enable_static()