【HEU】[Paddle Tensor 第二期 API鲁棒性增强] paddle.roll 支持 0-size tensor 与 bool (#70537)

fangfangssj · web-flow · commit 80f588a71802 · 2025-01-06T10:59:22.000+08:00
diff --git a/paddle/phi/kernels/cpu/roll_kernel.cc b/paddle/phi/kernels/cpu/roll_kernel.cc
@@ -18,6 +18,7 @@
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/cast_kernel.h"
 #include "paddle/phi/kernels/cpu/roll_kernel_impl.h"
 
 namespace phi {
@@ -28,8 +29,21 @@ void RollKernel(const Context& dev_ctx,
                 const IntArray& shifts,
                 const std::vector<int64_t>& axis,
                 DenseTensor* out) {
-  std::vector<T> out_vec;
-  phi::TensorToVector(x, dev_ctx, &out_vec);
+  if (x.numel() == 0) {
+    out->Resize(out->dims());
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
+  using Type =
+      typename std::conditional<std::is_same<T, bool>::value, int16_t, T>::type;
+  std::vector<Type> out_vec;
+  if (std::is_same<T, bool>::value) {
+    DenseTensor tmp_int_tensor;
+    tmp_int_tensor = phi::Cast<T, Context>(dev_ctx, x, phi::DataType::INT16);
+    phi::TensorToVector(tmp_int_tensor, dev_ctx, &out_vec);
+  } else {
+    phi::TensorToVector(x, dev_ctx, &out_vec);
+  }
 
   auto shifts_data = shifts.GetData();
   size_t nums = shifts_data.size();
@@ -57,7 +71,14 @@ void RollKernel(const Context& dev_ctx,
     ShiftAlongDim(out_vec.data(), input_dim, dims[i], shifts_data[i]);
   }
   dev_ctx.template Alloc<T>(out);
-  phi::TensorFromVector(out_vec, dev_ctx, out);
+  if (std::is_same<T, bool>::value) {
+    DenseTensor tmp_bool_tensor;
+    phi::TensorFromVector(out_vec, dev_ctx, &tmp_bool_tensor);
+    *out =
+        phi::Cast<Type, Context>(dev_ctx, tmp_bool_tensor, phi::DataType::BOOL);
+  } else {
+    phi::TensorFromVector(out_vec, dev_ctx, out);
+  }
   out->Resize(x.dims());
 }
 
@@ -67,6 +88,7 @@ PD_REGISTER_KERNEL(roll,
                    CPU,
                    ALL_LAYOUT,
                    phi::RollKernel,
+                   bool,
                    float,
                    double,
                    int,
diff --git a/paddle/phi/kernels/gpu/roll_kernel.cu b/paddle/phi/kernels/gpu/roll_kernel.cu
@@ -29,6 +29,11 @@ void RollKernel(const Context& dev_ctx,
                 const IntArray& shifts,
                 const std::vector<int64_t>& axis,
                 DenseTensor* out) {
+  if (x.numel() == 0) {
+    out->Resize(out->dims());
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   auto* in_data = x.data<T>();
   T* out_data = dev_ctx.template Alloc<T>(out);
 
@@ -69,6 +74,7 @@ PD_REGISTER_KERNEL(roll,
                    phi::RollKernel,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
+                   bool,
                    float,
                    double,
                    int,
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -2106,6 +2106,7 @@ def roll(
             x,
             'dtype',
             [
+                'bool',
                 'float16',
                 'float32',
                 'uint16',
diff --git a/test/legacy_test/test_roll_op.py b/test/legacy_test/test_roll_op.py

Original file line number	Diff line number	Diff line change
`@@ -2106,6 +2106,7 @@ def roll(`
`2106`	`2106`	`x,`
`2107`	`2107`	`'dtype',`
`2108`	`2108`	`[`
	`2109`	`+ 'bool',`
`2109`	`2110`	`'float16',`
`2110`	`2111`	`'float32',`
`2111`	`2112`	`'uint16',`