From f5e430c53e31b0c03c057d9318c46cdfb075d195 Mon Sep 17 00:00:00 2001
From: WeiXin <weixin10@baidu.com>
Date: Tue, 14 Sep 2021 10:21:32 +0800
Subject: [PATCH] slice_op support bool tensor. (#35586)

---
 paddle/fluid/operators/eigen/pad.cc              |  1 +
 paddle/fluid/operators/eigen/pad.cu              |  1 +
 paddle/fluid/operators/eigen/slice.cu            |  1 +
 paddle/fluid/operators/slice_op.cc               | 11 +++++++----
 .../fluid/tests/unittests/test_slice_op.py       | 16 ++++++++++++++++
 5 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/paddle/fluid/operators/eigen/pad.cc b/paddle/fluid/operators/eigen/pad.cc
index 421c9eaf5cd..9db4571357a 100644
--- a/paddle/fluid/operators/eigen/pad.cc
+++ b/paddle/fluid/operators/eigen/pad.cc
@@ -51,6 +51,7 @@ struct EigenPad<Eigen::DefaultDevice, T, Rank> {
   template struct FUNCTOR<Eigen::DefaultDevice, TYPE, 4>; \
   template struct FUNCTOR<Eigen::DefaultDevice, TYPE, 5>; \
   template struct FUNCTOR<Eigen::DefaultDevice, TYPE, 6>
+INSTANTIATION(EigenPad, bool);
 INSTANTIATION(EigenPad, int);
 INSTANTIATION(EigenPad, int64_t);
 INSTANTIATION(EigenPad, float);
diff --git a/paddle/fluid/operators/eigen/pad.cu b/paddle/fluid/operators/eigen/pad.cu
index 4cf88712d95..e028a8aef18 100644
--- a/paddle/fluid/operators/eigen/pad.cu
+++ b/paddle/fluid/operators/eigen/pad.cu
@@ -53,6 +53,7 @@ struct EigenPad<Eigen::GpuDevice, T, Rank> {
   template struct FUNCTOR<Eigen::GpuDevice, TYPE, 4>; \
   template struct FUNCTOR<Eigen::GpuDevice, TYPE, 5>; \
   template struct FUNCTOR<Eigen::GpuDevice, TYPE, 6>
+INSTANTIATION(EigenPad, bool);
 INSTANTIATION(EigenPad, int);
 INSTANTIATION(EigenPad, int64_t);
 INSTANTIATION(EigenPad, float);
diff --git a/paddle/fluid/operators/eigen/slice.cu b/paddle/fluid/operators/eigen/slice.cu
index dc51fa72220..3dfd0500cc9 100644
--- a/paddle/fluid/operators/eigen/slice.cu
+++ b/paddle/fluid/operators/eigen/slice.cu
@@ -53,6 +53,7 @@ struct EigenSlice<Eigen::GpuDevice, T, Rank> {
   template struct FUNCTOR<Eigen::GpuDevice, TYPE, 4>; \
   template struct FUNCTOR<Eigen::GpuDevice, TYPE, 5>; \
   template struct FUNCTOR<Eigen::GpuDevice, TYPE, 6>
+INSTANTIATION(EigenSlice, bool);
 INSTANTIATION(EigenSlice, int);
 INSTANTIATION(EigenSlice, int64_t);
 INSTANTIATION(EigenSlice, float);
diff --git a/paddle/fluid/operators/slice_op.cc b/paddle/fluid/operators/slice_op.cc
index 2a13998b4f6..a5513ba6487 100644
--- a/paddle/fluid/operators/slice_op.cc
+++ b/paddle/fluid/operators/slice_op.cc
@@ -434,7 +434,8 @@ REGISTER_OPERATOR(slice_grad, ops::SliceOpGrad,
                   ops::SliceOpGradVarTypeInference);
 
 REGISTER_OP_CPU_KERNEL(
-    slice, ops::SliceKernel<paddle::platform::CPUDeviceContext, int>,
+    slice, ops::SliceKernel<paddle::platform::CPUDeviceContext, bool>,
+    ops::SliceKernel<paddle::platform::CPUDeviceContext, int>,
     ops::SliceKernel<paddle::platform::CPUDeviceContext, int64_t>,
     ops::SliceKernel<paddle::platform::CPUDeviceContext, float>,
     ops::SliceKernel<paddle::platform::CPUDeviceContext, double>,
@@ -444,7 +445,8 @@ REGISTER_OP_CPU_KERNEL(
                      paddle::platform::complex<double>>);
 
 REGISTER_OP_CPU_KERNEL(
-    slice_grad, ops::SliceGradKernel<paddle::platform::CPUDeviceContext, int>,
+    slice_grad, ops::SliceGradKernel<paddle::platform::CPUDeviceContext, bool>,
+    ops::SliceGradKernel<paddle::platform::CPUDeviceContext, int>,
     ops::SliceGradKernel<paddle::platform::CPUDeviceContext, int64_t>,
     ops::SliceGradKernel<paddle::platform::CPUDeviceContext, float>,
     ops::SliceGradKernel<paddle::platform::CPUDeviceContext, double>,
@@ -454,7 +456,8 @@ REGISTER_OP_CPU_KERNEL(
                          paddle::platform::complex<double>>);
 
 REGISTER_OP_CUDA_KERNEL(
-    slice, ops::SliceKernel<paddle::platform::CUDADeviceContext, float>,
+    slice, ops::SliceKernel<paddle::platform::CUDADeviceContext, bool>,
+    ops::SliceKernel<paddle::platform::CUDADeviceContext, float>,
     ops::SliceKernel<paddle::platform::CUDADeviceContext, double>,
     ops::SliceKernel<paddle::platform::CUDADeviceContext, int>,
     ops::SliceKernel<paddle::platform::CUDADeviceContext, int64_t>,
@@ -466,7 +469,7 @@ REGISTER_OP_CUDA_KERNEL(
                      paddle::platform::complex<double>>);
 
 REGISTER_OP_CUDA_KERNEL(
-    slice_grad,
+    slice_grad, ops::SliceGradKernel<paddle::platform::CUDADeviceContext, bool>,
     ops::SliceGradKernel<paddle::platform::CUDADeviceContext, float>,
     ops::SliceGradKernel<paddle::platform::CUDADeviceContext, double>,
     ops::SliceGradKernel<paddle::platform::CUDADeviceContext, int>,
diff --git a/python/paddle/fluid/tests/unittests/test_slice_op.py b/python/paddle/fluid/tests/unittests/test_slice_op.py
index a80dc87525a..57d5453ec96 100644
--- a/python/paddle/fluid/tests/unittests/test_slice_op.py
+++ b/python/paddle/fluid/tests/unittests/test_slice_op.py
@@ -553,6 +553,22 @@ class TestSliceApiWithTensor(unittest.TestCase):
 
             self.assertTrue(np.array_equal(a_1.numpy(), a_2.numpy()))
 
+    def test_bool_tensor(self):
+        with paddle.fluid.dygraph.guard():
+            array = (np.arange(60).reshape([3, 4, 5]) % 3).astype('bool')
+            tt = paddle.to_tensor(array)
+            tt.stop_gradient = False
+
+            starts = [0, 1, 2]
+            ends = [3, 5, 4]
+            axes = [0, 1, 2]
+
+            y_paddle = paddle.slice(tt, axes, starts, ends)
+            y_np = tt[0:3, 1:5, 2:4]
+
+            self.assertTrue(paddle.bool == y_paddle.dtype)
+            self.assertTrue(np.array_equal(y_paddle.numpy(), y_np))
+
 
 class TestSliceApiWithLoDTensorArray(unittest.TestCase):
     def setUp(self):
-- 
GitLab