【Paddle Tensor 第二期 API支持 0-size Tensor】paddle.full_like 支持 0-size tensor-part (#70077)

SCUcookie · web-flow · commit 15aa29edb22f · 2024-12-13T10:46:54.000+08:00
* submit

* modify

* modify test
diff --git a/paddle/phi/kernels/cpu/full_kernel.cc b/paddle/phi/kernels/cpu/full_kernel.cc
@@ -25,6 +25,9 @@ namespace phi {
 template <typename T, typename Context, typename VType>
 void FullValue(const Context& dev_ctx, DenseTensor* tensor, VType val) {
   dev_ctx.template Alloc<T>(tensor);
+  if (tensor->numel() == 0) {
+    return;
+  }
   auto t = phi::EigenVector<T>::Flatten(*tensor);
   t.device(*dev_ctx.eigen_device()) = t.constant(static_cast<T>(val));
 }
@@ -36,6 +39,10 @@ void FullKernel(const Context& dev_ctx,
                 DataType dtype UNUSED,
                 DenseTensor* out) {
   out->Resize(common::make_ddim(shape.GetData()));
+  if (out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   FullValue<T>(dev_ctx, out, val.to<T>());
 }
 
@@ -45,6 +52,11 @@ void FullLikeKernel(const Context& dev_ctx,
                     const Scalar& val,
                     DataType dtype UNUSED,
                     DenseTensor* out) {
+  if (out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    out->Resize(x.dims());
+    return;
+  }
   if (!std::is_same<T, phi::dtype::complex<float>>::value &&
       !std::is_same<T, phi::dtype::complex<double>>::value) {
     auto value = val.to<double>();
@@ -93,6 +105,9 @@ void FullIntArrayKernel(const Context& dev_ctx,
                         DenseTensor* out) {
   out->Resize(common::make_ddim({static_cast<int64_t>(shape.size())}));
   T* out_data = dev_ctx.template Alloc<T>(out);
+  if (out->numel() == 0) {
+    return;
+  }
   for (size_t i = 0; i < shape.size(); ++i) {
     int64_t val = shape[i];
     out_data[i] = static_cast<T>(val);
diff --git a/paddle/phi/kernels/gpu/full_kernel.cu b/paddle/phi/kernels/gpu/full_kernel.cu
@@ -46,7 +46,7 @@ void FullKernel(const Context& dev_ctx,
   dev_ctx.template Alloc<T>(out);
 
   if (numel > 0) {
-    // in transformer model the numel of outpout will be zero.
+    // in transformer model the numel of output will be zero.
     std::vector<const DenseTensor*> inputs = {};
     std::vector<DenseTensor*> outputs = {out};
     // This function has no input, so the inputs.size() == 0. Use kUnary, but
diff --git a/paddle/phi/kernels/xpu/full_kernel.cc b/paddle/phi/kernels/xpu/full_kernel.cc
@@ -36,8 +36,8 @@ void FullKernel(const Context& dev_ctx,
   using XPUInTDType = typename XPUTypeTrait<T>::Type;
   out->Resize(common::make_ddim(shape.GetData()));
   dev_ctx.template Alloc<T>(out);
-  auto out_data = reinterpret_cast<XPUInTDType*>(out->data<T>());
   if (out->numel() > 0) {
+    auto out_data = reinterpret_cast<XPUInTDType*>(out->data<T>());
     int r = xpu::constant(dev_ctx.x_context(),
                           out_data,
                           out->numel(),
@@ -53,40 +53,40 @@ void FullLikeKernel(const Context& dev_ctx,
                     DataType dtype,
                     DenseTensor* out) {
   dev_ctx.template Alloc<T>(out);
-  auto value = val.to<double>();
-  using XPUInTDType = typename XPUTypeTrait<T>::Type;
-  using CommonType = typename std::common_type<
-      float,
-      typename std::conditional<std::is_same<T, phi::dtype::float16>::value,
-                                float,
-                                T>::type>::type;
+  if (out->numel() > 0) {
+    auto value = val.to<double>();
+    using XPUInTDType = typename XPUTypeTrait<T>::Type;
+    using CommonType = typename std::common_type<
+        float,
+        typename std::conditional<std::is_same<T, phi::dtype::float16>::value,
+                                  float,
+                                  T>::type>::type;
 
-  auto common_type_value = static_cast<CommonType>(value);
-  bool is_out_range = true;
-  if (std::isinf(value) || std::isnan(value)) {
-    is_out_range = false;
-  }
-  if ((common_type_value >=
-       static_cast<CommonType>(std::numeric_limits<T>::lowest())) &&
-      (common_type_value <=
-       static_cast<CommonType>(std::numeric_limits<T>::max()))) {
-    is_out_range = false;
-  }
+    auto common_type_value = static_cast<CommonType>(value);
+    bool is_out_range = true;
+    if (std::isinf(value) || std::isnan(value)) {
+      is_out_range = false;
+    }
+    if ((common_type_value >=
+         static_cast<CommonType>(std::numeric_limits<T>::lowest())) &&
+        (common_type_value <=
+         static_cast<CommonType>(std::numeric_limits<T>::max()))) {
+      is_out_range = false;
+    }
 
-  PADDLE_ENFORCE_EQ(
-      is_out_range,
-      false,
-      common::errors::InvalidArgument(
-          "The filled value is out of range for target type, "
-          "current kernel type is %s, the range should between %f "
-          "and %f, but now value is %f.",
-          typeid(T).name(),
-          static_cast<CommonType>(std::numeric_limits<T>::lowest()),
-          static_cast<CommonType>(std::numeric_limits<T>::max()),
-          static_cast<float>(value)));
+    PADDLE_ENFORCE_EQ(
+        is_out_range,
+        false,
+        common::errors::InvalidArgument(
+            "The filled value is out of range for target type, "
+            "current kernel type is %s, the range should between %f "
+            "and %f, but now value is %f.",
+            typeid(T).name(),
+            static_cast<CommonType>(std::numeric_limits<T>::lowest()),
+            static_cast<CommonType>(std::numeric_limits<T>::max()),
+            static_cast<float>(value)));
 
-  auto out_data = reinterpret_cast<XPUInTDType*>(out->data<T>());
-  if (out->numel() > 0) {
+    auto out_data = reinterpret_cast<XPUInTDType*>(out->data<T>());
     int r = xpu::constant(dev_ctx.x_context(),
                           out_data,
                           out->numel(),
diff --git a/test/legacy_test/test_full_like_op.py b/test/legacy_test/test_full_like_op.py
@@ -218,5 +218,56 @@ def init_data(self):
         self.dtype = np.uint16
 
 
+class TestFullKernelZeroSize(unittest.TestCase):
+    def test_full_kernel_cpu_zero_size(self):
+        paddle.disable_static()
+        value = 5
+        dtype = "int32"
+        shape = [0, 3]
+        tensor = paddle.full(shape, value, dtype=dtype)
+        expected = np.full(shape, value, dtype=dtype)
+        self.assertTrue(np.array_equal(tensor.numpy(), expected))
+        paddle.enable_static()
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda(), "Paddle is not compiled with CUDA"
+    )
+    def test_full_kernel_gpu_zero_size(self):
+        paddle.disable_static()
+        paddle.set_device("gpu:0")
+        value = 5.5
+        dtype = "float32"
+        shape = [0, 3]
+        tensor = paddle.full(shape, value, dtype=dtype)
+        expected = np.full(shape, value, dtype=dtype)
+        self.assertTrue(np.array_equal(tensor.numpy(), expected))
+        paddle.enable_static()
+
+
+class TestFullLikeKernelZeroSize(unittest.TestCase):
+    def test_full_like_kernel_cpu_zero_size(self):
+        paddle.disable_static()
+        base_tensor = paddle.to_tensor(np.empty((0, 2), dtype=np.float32))
+        value = 10.0
+        result = paddle.full_like(base_tensor, value, dtype="float32")
+        expected = np.full_like(base_tensor.numpy(), value)
+        self.assertTrue(np.array_equal(result.numpy(), expected))
+        paddle.enable_static()
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda(), "Paddle is not compiled with CUDA"
+    )
+    def test_full_like_kernel_gpu_zero_size(self):
+        paddle.disable_static()
+        base_tensor = paddle.to_tensor(
+            np.empty((0, 3), dtype=np.float32), place=paddle.CUDAPlace(0)
+        )
+        value = 20.0
+        result = paddle.full_like(base_tensor, value, dtype="float32")
+        expected = np.full_like(base_tensor.numpy(), value)
+        self.assertTrue(np.array_equal(result.numpy(), expected))
+        paddle.enable_static()
+
+
 if __name__ == "__main__":
     unittest.main()
