[0-size Tensor No.174] Add 0-size Tensor support for grid_sample (#72972)

co63oc · web-flow · commit 6c675da309b0 · 2025-05-31T20:16:11.000+08:00
* Fix

* Fix
diff --git a/paddle/phi/kernels/cpu/grid_sample_grad_kernel.cc b/paddle/phi/kernels/cpu/grid_sample_grad_kernel.cc
@@ -16,9 +16,9 @@
 
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cpu/grid_sample_utils.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
-
 namespace phi {
 
 template <typename T>
@@ -653,12 +653,26 @@ template <typename T, typename Context>
 void GridSampleGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           const DenseTensor& grid,
-                          const DenseTensor& out_grid,
+                          const DenseTensor& out_grad,
                           const std::string& mode,
                           const std::string& padding_mode,
                           bool align_corners,
                           DenseTensor* x_grad,
                           DenseTensor* grid_grad) {
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (grid_grad) {
+      phi::Full<T, Context>(dev_ctx,
+                            phi::IntArray(common::vectorize(grid_grad->dims())),
+                            0,
+                            grid_grad);
+    }
+    return;
+  }
+
   if (x.dims().size() == 4) {
     const int n = static_cast<int>(grid.dims()[0]);
     const int out_h = static_cast<int>(grid.dims()[1]);
@@ -693,7 +707,7 @@ void GridSampleGradKernel(const Context& dev_ctx,
     if (mode == "bilinear") {
       GatherBilinearGrad<T>(dev_ctx,
                             x,
-                            out_grid,
+                            out_grad,
                             &grid_x,
                             &grid_y,
                             &grid_x_scale,
@@ -705,7 +719,7 @@ void GridSampleGradKernel(const Context& dev_ctx,
       auto grid_y_t = EigenTensor<T, 3>::From(grid_y);
       grid_x_t = grid_x_t.round();
       grid_y_t = grid_y_t.round();
-      GatherOutputGradToInputGrad<T>(out_grid, x_grad, grid_x, grid_y);
+      GatherOutputGradToInputGrad<T>(out_grad, x_grad, grid_x, grid_y);
     }
   } else {
     const int n = static_cast<int>(grid.dims()[0]);
@@ -746,7 +760,7 @@ void GridSampleGradKernel(const Context& dev_ctx,
     if (mode == "bilinear") {
       Gather3DBilinearGrad<T>(dev_ctx,
                               x,
-                              out_grid,
+                              out_grad,
                               &grid_x,
                               &grid_y,
                               &grid_z,
@@ -757,7 +771,7 @@ void GridSampleGradKernel(const Context& dev_ctx,
                               grid_grad);
     } else {
       Gather3DOutputGradToInputGrad<T>(
-          out_grid, x_grad, grid_x, grid_y, grid_z);
+          out_grad, x_grad, grid_x, grid_y, grid_z);
     }
   }
 }
diff --git a/paddle/phi/kernels/cpu/grid_sample_kernel.cc b/paddle/phi/kernels/cpu/grid_sample_kernel.cc
@@ -312,6 +312,10 @@ void GridSampleKernel(const Context& dev_ctx,
                       const std::string& padding_mode,
                       bool align_corners,
                       DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   if (x.dims().size() == 4) {
     const int n = static_cast<int>(grid.dims()[0]);
     const int out_h = static_cast<int>(grid.dims()[1]);
diff --git a/paddle/phi/kernels/gpu/grid_sample_grad_kernel.cu b/paddle/phi/kernels/gpu/grid_sample_grad_kernel.cu
@@ -19,6 +19,7 @@
 #include "paddle/phi/backends/gpu/gpu_launch_config.h"
 #include "paddle/phi/backends/gpu/gpu_primitives.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/gpu/grid_sample_utils.h"
 
@@ -572,6 +573,20 @@ void GridSampleGradKernel(const Context& dev_ctx,
                           bool align_corners,
                           DenseTensor* x_grad,
                           DenseTensor* grid_grad) {
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (grid_grad) {
+      phi::Full<T, Context>(dev_ctx,
+                            phi::IntArray(common::vectorize(grid_grad->dims())),
+                            0,
+                            grid_grad);
+    }
+    return;
+  }
+
   PaddingMode enum_padding_mode;
   Mode enum_mode;
   if (padding_mode == "border") {
diff --git a/paddle/phi/kernels/gpu/grid_sample_kernel.cu b/paddle/phi/kernels/gpu/grid_sample_kernel.cu
@@ -323,6 +323,10 @@ void GridSampleKernel(const Context& dev_ctx,
                       const std::string& padding_mode,
                       bool align_corners,
                       DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   PaddingMode enum_padding_mode;
   Mode enum_mode;
   if (padding_mode == "border") {
diff --git a/paddle/phi/kernels/xpu/grid_sample_grad_kernel.cc b/paddle/phi/kernels/xpu/grid_sample_grad_kernel.cc
@@ -16,14 +16,14 @@
 
 #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>
 void GridSampleGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           const DenseTensor& grid,
-                          const DenseTensor& out_grid,
+                          const DenseTensor& out_grad,
                           const std::string& mode,
                           const std::string& padding_mode,
                           bool align_corners,
@@ -35,6 +35,19 @@ void GridSampleGradKernel(const Context& dev_ctx,
       common::errors::InvalidArgument(
           ("XPU is only support input_dims == 4 in grid_sample_grad op.")));
 
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (grid_grad) {
+      phi::Full<T, Context>(dev_ctx,
+                            phi::IntArray(common::vectorize(grid_grad->dims())),
+                            0,
+                            grid_grad);
+    }
+    return;
+  }
   const int64_t n = grid.dims()[0];
   const int64_t out_h = grid.dims()[1];
   const int64_t out_w = grid.dims()[2];
@@ -65,7 +78,7 @@ void GridSampleGradKernel(const Context& dev_ctx,
   int r = xpu::grid_sample_grad<T>(dev_ctx.x_context(),
                                    x.data<T>(),
                                    grid.data<T>(),
-                                   out_grid.data<T>(),
+                                   out_grad.data<T>(),
                                    x_grad_ptr,
                                    grid_grad_ptr,
                                    n,
diff --git a/paddle/phi/kernels/xpu/grid_sample_kernel.cc b/paddle/phi/kernels/xpu/grid_sample_kernel.cc
@@ -28,6 +28,10 @@ void GridSampleKernel(const Context& dev_ctx,
                       const std::string& padding_mode,
                       bool align_corners,
                       DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   // attrs
   // paddle.nn.functional.grid_sample(x, grid, mode='bilinear',
   // padding_mode='zeros', align_corners=True, name=None)
diff --git a/test/legacy_test/test_grid_sampler_op.py b/test/legacy_test/test_grid_sampler_op.py
@@ -451,6 +451,16 @@ def initTestCase(self):
         self.numeric_grad_delta = 0.0001
 
 
+class Case_ZeroSize(TestGridSamplerOp):
+    def initTestCase(self):
+        self.x_shape = (2, 0, 5, 6)
+        self.grid_shape = (2, 8, 9, 2)
+        self.theta_shape = (2, 2, 3)
+        self.align_corners = False
+        self.padding_mode = "zeros"
+        self.mode = "bilinear"
+
+
 @skip_check_grad_ci(
     reason="'check_grad' on large inputs is too slow, "
     + "however it is desirable to cover the forward pass"
