【Paddle Tensor 第二期 API支持 0-size Tensor】paddle.linspace 支持 0-size tensor (#70047)

aquagull · web-flow · commit d785cf4bfffb · 2024-12-11T16:22:16.000+08:00
* compatible to 0-size

* add test

* fix

* refine
diff --git a/paddle/phi/kernels/cpu/linspace_kernel.cc b/paddle/phi/kernels/cpu/linspace_kernel.cc
@@ -33,17 +33,22 @@ void LinspaceKernel(const Context& ctx,
   } else if (number.dtype() == phi::DataType::INT32) {
     num = number.data<int32_t>()[0];
   }
+  PADDLE_ENFORCE_GE(num,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The num of linspace op should be larger "
+                        "than or equal to 0, but received num is %d",
+                        num));
+  if (num == 0) {
+    out->Resize(common::make_ddim({0}));
+    ctx.template Alloc<T>(out);
+    return;
+  }
   auto start_t = phi::funcs::TransDataType(ctx, start, dtype);
   auto stop_t = phi::funcs::TransDataType(ctx, stop, dtype);
 
   T start_data = start_t.template data<T>()[0];
   T stop_data = stop_t.template data<T>()[0];
-  PADDLE_ENFORCE_GT(
-      num,
-      0,
-      common::errors::InvalidArgument("The num of linspace op should be larger "
-                                      "than 0, but received num is %d",
-                                      num));
 
   out->Resize(common::make_ddim({num}));
   T* out_data = ctx.template Alloc<T>(out);
diff --git a/paddle/phi/kernels/gpu/linspace_kernel.cu b/paddle/phi/kernels/gpu/linspace_kernel.cu
@@ -80,17 +80,18 @@ void LinspaceKernel(const Context& ctx,
   T start_value = GetValueOfExpectedType<T, Context>(ctx, start);
   T stop_value = GetValueOfExpectedType<T, Context>(ctx, stop);
   int64_t num = GetValueOfExpectedType<int64_t, Context>(ctx, number);
-
-  PADDLE_ENFORCE_GT(
-      num,
-      0,
-      common::errors::InvalidArgument("The num of linspace op should be larger "
-                                      "than 0, but received num is %d",
-                                      num));
+  PADDLE_ENFORCE_GE(num,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The num of linspace op should be larger "
+                        "than or equal to 0, but received num is %d",
+                        num));
 
   out->Resize(common::make_ddim({num}));
   T* out_data = ctx.template Alloc<T>(out);
-
+  if (num == 0) {
+    return;
+  }
   auto stream = ctx.stream();
   if (num != 1) {
     int block = 512;
diff --git a/paddle/phi/kernels/xpu/linspace_kernel.cc b/paddle/phi/kernels/xpu/linspace_kernel.cc
@@ -59,17 +59,18 @@ void LinspaceKernel(const Context& ctx,
   T start_value = GetValueOfExpectedType<T, Context>(ctx, start);
   T stop_value = GetValueOfExpectedType<T, Context>(ctx, stop);
   int64_t num = GetValueOfExpectedType<int64_t, Context>(ctx, number);
-
-  PADDLE_ENFORCE_GT(
-      num,
-      0,
-      common::errors::InvalidArgument("The num of linspace op should be larger "
-                                      "than 0, but received num is %d",
-                                      num));
+  PADDLE_ENFORCE_GE(num,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The num of linspace op should be larger "
+                        "than or equal to 0, but received num is %d",
+                        num));
 
   out->Resize(common::make_ddim({num}));
   T* out_data = ctx.template Alloc<T>(out);
-
+  if (num == 0) {
+    return;
+  }
   int r = xpu::linspace(ctx.x_context(),
                         reinterpret_cast<XPUType*>(out_data),
                         static_cast<XPUType>(start_value),
diff --git a/test/legacy_test/test_linspace.py b/test/legacy_test/test_linspace.py
@@ -16,6 +16,7 @@
 
 import numpy as np
 from op_test import OpTest, convert_float_to_uint16, paddle_static_guard
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle import base
@@ -246,5 +247,37 @@ def test_num_dtype():
                 self.assertRaises(TypeError, test_step_dtype)
 
 
+class TestLinspaceOpEmptyTensor(unittest.TestCase):
+    def _get_places(self):
+        places = [base.CPUPlace()]
+        if paddle.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        return places
+
+    def _test_linspace_empty_static(self, place):
+        with static_guard():
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                out = paddle.linspace(0, 10, 0, dtype='float32')
+                exe = paddle.static.Executor(place)
+                res = exe.run(fetch_list=[out])
+                self.assertEqual(res[0].shape, (0,))
+                self.assertEqual(len(res[0]), 0)
+
+    def _test_linspace_empty_dynamic(self):
+        with dygraph_guard():
+            out = paddle.linspace(0, 10, 0, dtype='float32')
+            self.assertEqual(out.shape, [0])
+            self.assertEqual(len(out.numpy()), 0)
+
+    def test_empty_tensor(self):
+        places = self._get_places()
+        for place in places:
+            self._test_linspace_empty_static(place)
+
+        self._test_linspace_empty_dynamic()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/xpu/test_linspace_op_xpu.py b/test/xpu/test_linspace_op_xpu.py
@@ -79,6 +79,15 @@ def set_attrs(self):
             }
             self.outputs = {'Out': np.array(10, dtype=self.dtype)}
 
+    class TestXPULinespace4(TestXPULinespaceOp):
+        def set_attrs(self):
+            self.inputs = {
+                'Start': np.array([0]).astype(self.dtype),
+                'Stop': np.array([10]).astype(self.dtype),
+                'Num': np.array([0]).astype('int32'),
+            }
+            self.outputs = {'Out': np.array([], dtype=self.dtype)}
+
 
 support_types = get_xpu_op_support_types('linspace')
 for stype in support_types:
