[0-size Tensor No.101、311] Add 0-size Tensor support for lerp (#72906)

co63oc · web-flow · commit 892456df95fa · 2025-05-28T10:56:54.000+08:00
* Fix

* Fix

* Fix

* Fix

* Fix

* Fix
diff --git a/paddle/phi/kernels/gpu/lerp_grad_kernel.cu b/paddle/phi/kernels/gpu/lerp_grad_kernel.cu
@@ -21,6 +21,7 @@
 #include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/kernels/broadcast_tensors_kernel.h"
 #include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
@@ -174,6 +175,18 @@ void LerpGradKernel(const Context& ctx,
                     const DenseTensor& out_grad,
                     DenseTensor* x_grad,
                     DenseTensor* y_grad) {
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(
+          ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (y_grad) {
+      phi::Full<T, Context>(
+          ctx, phi::IntArray(common::vectorize(y_grad->dims())), 0, y_grad);
+    }
+    return;
+  }
+
   const int rank = out.dims().size();
   PADDLE_ENFORCE_GE(
       rank,
diff --git a/paddle/phi/kernels/gpu/lerp_kernel.cu b/paddle/phi/kernels/gpu/lerp_kernel.cu
@@ -47,15 +47,10 @@ void LerpKernel(const Context &ctx,
                 const DenseTensor &y,
                 const DenseTensor &weight,
                 DenseTensor *out) {
-  PADDLE_ENFORCE_GT(
-      x.numel(),
-      0,
-      common::errors::InvalidArgument("LerpKernel's input x must not empty."));
-
-  PADDLE_ENFORCE_GT(
-      y.numel(),
-      0,
-      common::errors::InvalidArgument("LerpKernel's input y must not empty."));
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
 
   int rank = out->dims().size();
   PADDLE_ENFORCE_GE(
diff --git a/paddle/phi/kernels/impl/lerp_grad_kernel_impl.h b/paddle/phi/kernels/impl/lerp_grad_kernel_impl.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 
@@ -28,6 +29,18 @@ static void LerpGradFunction(const Context& ctx,
                              const DenseTensor& out_grad,
                              DenseTensor* x_grad,
                              DenseTensor* y_grad) {
+  if (out_grad.numel() == 0) {
+    if (x_grad) {
+      phi::Full<T, Context>(
+          ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (y_grad) {
+      phi::Full<T, Context>(
+          ctx, phi::IntArray(common::vectorize(y_grad->dims())), 0, y_grad);
+    }
+    return;
+  }
+
   auto& w = weight;
   auto& dout = out_grad;
   auto* dx = x_grad;
diff --git a/paddle/phi/kernels/impl/lerp_kernel_impl.h b/paddle/phi/kernels/impl/lerp_kernel_impl.h
@@ -83,16 +83,10 @@ void LerpKernel(const Context& ctx,
                 const DenseTensor& y,
                 const DenseTensor& weight,
                 DenseTensor* out) {
-  PADDLE_ENFORCE_GT(
-      x.numel(),
-      0,
-      common::errors::InvalidArgument("LerpKernel's input x must not empty."));
-
-  PADDLE_ENFORCE_GT(
-      y.numel(),
-      0,
-      common::errors::InvalidArgument("LerpKernel's input y must not empty."));
-
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
   int rank = out->dims().size();
   PADDLE_ENFORCE_GE(
       rank,
diff --git a/test/legacy_test/test_lerp_op.py b/test/legacy_test/test_lerp_op.py
@@ -35,8 +35,12 @@ def setUp(self):
         self.init_shape()
         self.init_xyshape()
         self.init_wshape()
-        x = np.arange(1.0, 101.0).astype(self.dtype).reshape(self.xshape)
-        y = np.full(100, 10.0).astype(self.dtype).reshape(self.yshape)
+        if 0 in self.shape:
+            x = np.random.rand(*self.xshape).astype(self.dtype)
+            y = np.random.rand(*self.yshape).astype(self.dtype)
+        else:
+            x = np.arange(1.0, 101.0).astype(self.dtype).reshape(self.xshape)
+            y = np.full(100, 10.0).astype(self.dtype).reshape(self.yshape)
         w = np.random.random(self.wshape).astype(self.dtype)
         self.inputs = {'X': x, 'Y': y, 'Weight': w}
         self.outputs = {'Out': x + w * (y - x)}
@@ -94,6 +98,11 @@ def init_dtype(self):
         self.dtype = np.float16
 
 
+class TestLerp_ZeroSize(TestLerp):
+    def init_shape(self):
+        self.shape = [2, 0]
+
+
 class TestLerpWihFp16BroadXY(TestLerp):
     def init_xyshape(self):
         self.xshape = [2, 1, 2, 5, 5]
