From 982d01ee233482a8d37a69a49dc664f336c4e1ba Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Fri, 15 Jul 2022 10:41:23 +0800
Subject: [PATCH 01/16] add paddle.take api

---
 python/paddle/__init__.py                     |   2 +
 .../paddle/fluid/tests/unittests/test_take.py | 136 ++++++++++++++++++
 python/paddle/tensor/__init__.py              |   2 +
 python/paddle/tensor/math.py                  |  54 +++++++
 4 files changed, 194 insertions(+)
 create mode 100644 python/paddle/fluid/tests/unittests/test_take.py

diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
index 6e47f4f9eab43..61ca4a7768756 100755
--- a/python/paddle/__init__.py
+++ b/python/paddle/__init__.py
@@ -275,6 +275,7 @@
 from .tensor.math import outer  # noqa: F401
 from .tensor.math import heaviside  # noqa: F401
 from .tensor.math import frac  # noqa: F401
+from .tensor.math import take  # noqa: F401
 
 from .tensor.random import bernoulli  # noqa: F401
 from .tensor.random import poisson  # noqa: F401
@@ -645,4 +646,5 @@
     'put_along_axis',
     'heaviside',
     'tril_indices',
+    'take',
 ]
diff --git a/python/paddle/fluid/tests/unittests/test_take.py b/python/paddle/fluid/tests/unittests/test_take.py
new file mode 100644
index 0000000000000..fc6b87032cf83
--- /dev/null
+++ b/python/paddle/fluid/tests/unittests/test_take.py
@@ -0,0 +1,136 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+import numpy as np
+import paddle
+import paddle.fluid as fluid
+import paddle.fluid.core as core
+from paddle.fluid import Program, program_guard
+
+
+class TestTakeAPI(unittest.TestCase):
+
+    def set_dtype(self):
+        self.input_dtype = 'float64'
+        self.index_dtype = 'int64'
+
+    def setUp(self):
+        self.set_dtype()
+        self.place = fluid.CUDAPlace(
+            0) if core.is_compiled_with_cuda() else fluid.CPUPlace()
+        self.input_shape = [3, 4]
+        self.index_shape = [2, 3]
+        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
+            self.input_dtype)
+        self.index_np = np.arange(-4, 2).reshape(self.index_shape).astype(
+            self.index_dtype)
+
+    def test_static_graph(self):
+        paddle.enable_static()
+        startup_program = Program()
+        train_program = Program()
+        with program_guard(startup_program, train_program):
+            x = fluid.data(name='input',
+                           dtype=self.input_dtype,
+                           shape=self.input_shape)
+            index = fluid.data(name='index',
+                               dtype=self.index_dtype,
+                               shape=self.index_shape)
+            out = paddle.take(x, index)
+
+            exe = fluid.Executor(self.place)
+            st_result = exe.run(fluid.default_main_program(),
+                                feed={
+                                    'input': self.input_np,
+                                    'index': self.index_np
+                                },
+                                fetch_list=[out])
+            self.assertTrue(
+                np.allclose(st_result, np.take(self.input_np, self.index_np)))
+
+    def test_dygraph(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.input_np)
+        index = paddle.to_tensor(self.index_np)
+        dy_result = paddle.take(x, index)
+        self.assertTrue(
+            np.allclose(np.take(self.input_np, self.index_np),
+                        dy_result.numpy()))
+
+
+class TestTakeInt32(TestTakeAPI):
+    """Test take API with data type int32"""
+
+    def set_dtype(self):
+        self.input_dtype = 'int32'
+        self.index_dtype = 'int64'
+
+
+class TestTakeInt64(TestTakeAPI):
+    """Test take API with data type int64"""
+
+    def set_dtype(self):
+        self.input_dtype = 'int64'
+        self.index_dtype = 'int64'
+
+
+class TestTakeFloat32(TestTakeAPI):
+    """Test take API with data type float32"""
+
+    def set_dtype(self):
+        self.input_dtype = 'float32'
+        self.index_dtype = 'int64'
+
+
+class TestTakeType(TestTakeAPI):
+    """Test take Error"""
+
+    def test_static_type_error(self):
+        """Argument 'index' must be Tensor"""
+        paddle.enable_static()
+        with program_guard(Program()):
+            x = fluid.data(name='input',
+                           dtype=self.input_dtype,
+                           shape=self.input_shape)
+            self.assertRaises(TypeError, paddle.take, x, self.index_np)
+
+    def test_dygraph_type_error(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.input_np)
+        self.assertRaises(TypeError, paddle.take, x, self.index_np)
+
+    def test_static_dtype_error(self):
+        """Data type of argument 'index' must be in [paddle.int32, paddle.int64]"""
+        paddle.enable_static()
+        with program_guard(Program()):
+            x = fluid.data(name='input',
+                           dtype='float64',
+                           shape=self.input_shape)
+            index = fluid.data(name='index',
+                               dtype='float32',
+                               shape=self.index_shape)
+            self.assertRaises(TypeError, paddle.take, x, index)
+
+    def test_dygraph_dtype_error(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.input_np)
+        index = paddle.to_tensor(self.index_np, dtype='float32')
+        self.assertRaises(TypeError, paddle.take, x, index)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
index 08b0af26bd46e..d4a90a95ee160 100755
--- a/python/paddle/tensor/__init__.py
+++ b/python/paddle/tensor/__init__.py
@@ -232,6 +232,7 @@
 from .math import outer  # noqa: F401
 from .math import heaviside  # noqa: F401
 from .math import frac  # noqa: F401
+from .math import take  # noqa: F401
 
 from .random import multinomial  # noqa: F401
 from .random import standard_normal  # noqa: F401
@@ -501,6 +502,7 @@
     'put_along_axis_',
     'exponential_',
     'heaviside',
+    'take',
 ]
 
 #this list used in math_op_patch.py for magic_method bind
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 0b5cc5bf6491b..f68a2e546105e 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -22,6 +22,7 @@
 from paddle.common_ops_import import OpProtoHolder
 from paddle.common_ops_import import templatedoc
 from paddle.common_ops_import import dygraph_utils
+from pip import main
 
 from .manipulation import cast
 from .creation import _complex_to_real_dtype
@@ -4549,3 +4550,56 @@ def frac(x, name=None):
             helper.append_op(
                 type="trunc", inputs=inputs, attrs=attrs, outputs={"Out": y})
             return _elementwise_op(LayerHelper(op_type, **locals()))
+
+def take(input, index, name=None):
+    """
+    Returns a new tensor with the elements of input at the given indices. 
+    The input tensor is treated as if it were viewed as a 1-D tensor. 
+    The result takes the same shape as the indices.
+    
+    Args:
+        input (Tensor): An N-D Tensor, which data type should be int32, int64, float32, float64. 
+        index (Tensor): An N-D Tensor, which data type should be int32, int64.
+        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        Tensor: Tensor with the same shape as index, the data type is the same with input.
+
+    Examples:
+        .. code-block:: python
+
+            import paddle
+
+            x = paddle.arange(0, 12).reshape([3, 4])
+            idx = paddle.arange(4, 10).reshape([2, 3])
+
+            paddle.take(x, idx)
+            # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
+            #        [[4, 5, 6],
+            #         [7, 8, 9]])
+
+            x.take(idx)
+            # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
+            #        [[4, 5, 6],
+            #         [7, 8, 9]])
+    """
+
+    if paddle.in_dynamic_mode():
+        if not isinstance(index, (paddle.Tensor, Variable)):
+            raise TypeError(
+                "The type of 'index' must be Tensor, but got {}".format(type(index)))
+        if index.dtype not in [paddle.int32, paddle.int64]:
+            raise TypeError(
+                "The data type of 'index' must be one of ['int32', 'int64'], but got {}".format(
+                    index.dtype))
+    else:
+        check_variable_and_dtype(index, 'index', ['int32', 'int64'], 'take')
+
+    input_1d = input.flatten()
+    index_1d = index.flatten()
+
+    # This processing enables 'take' to handle negative indexes within the correct range
+    index_1d = paddle.where(index_1d < 0, index_1d + input_1d.shape[0], index_1d)
+    out = input_1d.flatten().index_select(index_1d).reshape(index.shape)
+
+    return out

From b07c0624c1cc576bef7f66ac2213e42831a7b41b Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Fri, 29 Jul 2022 20:44:37 +0800
Subject: [PATCH 02/16] fix paddle.take

---
 python/paddle/tensor/math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index f68a2e546105e..0258ba7fe9924 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4600,6 +4600,6 @@ def take(input, index, name=None):
 
     # This processing enables 'take' to handle negative indexes within the correct range
     index_1d = paddle.where(index_1d < 0, index_1d + input_1d.shape[0], index_1d)
-    out = input_1d.flatten().index_select(index_1d).reshape(index.shape)
+    out = input_1d.index_select(index_1d).reshape(index.shape)
 
     return out

From c8482f62dc6e07f871526973a7c103e9d893dd83 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Fri, 29 Jul 2022 21:12:11 +0800
Subject: [PATCH 03/16] remove from pip import main

---
 python/paddle/tensor/math.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 6a71593286c77..9dc291baf0520 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -22,7 +22,6 @@
 from paddle.common_ops_import import OpProtoHolder
 from paddle.common_ops_import import templatedoc
 from paddle.common_ops_import import dygraph_utils
-from pip import main
 
 from .manipulation import cast
 from .creation import _complex_to_real_dtype

From 0665e5014506d07e927ac31934252be3f8c8f872 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Thu, 4 Aug 2022 20:28:50 +0800
Subject: [PATCH 04/16] test index out of range error

---
 .../paddle/fluid/tests/unittests/test_take.py | 38 ++++++++++++++++++-
 python/paddle/tensor/math.py                  | 11 +++---
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/python/paddle/fluid/tests/unittests/test_take.py b/python/paddle/fluid/tests/unittests/test_take.py
index fc6b87032cf83..a003f72def9eb 100644
--- a/python/paddle/fluid/tests/unittests/test_take.py
+++ b/python/paddle/fluid/tests/unittests/test_take.py
@@ -96,8 +96,8 @@ def set_dtype(self):
         self.index_dtype = 'int64'
 
 
-class TestTakeType(TestTakeAPI):
-    """Test take Error"""
+class TestTakeTypeError(TestTakeAPI):
+    """Test take Type Error"""
 
     def test_static_type_error(self):
         """Argument 'index' must be Tensor"""
@@ -132,5 +132,39 @@ def test_dygraph_dtype_error(self):
         self.assertRaises(TypeError, paddle.take, x, index)
 
 
+class TestTakeIndexRangeError(TestTakeAPI):
+    """Test take index out of range error"""
+
+    def setUp(self):
+        self.set_dtype()
+        self.place = fluid.CUDAPlace(
+            0) if core.is_compiled_with_cuda() else fluid.CPUPlace()
+        self.input_shape = [3, 4]
+        self.index_shape = [2, 3]
+        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
+            self.input_dtype)
+        self.index_np = np.arange(6, 12).reshape(self.index_shape).astype(
+            self.index_dtype)
+
+    def test_static_index_error(self):
+        """When the index is out of range,
+        an error is reported directly through `paddle.index_select`"""
+        paddle.enable_static()
+        with program_guard(Program()):
+            x = fluid.data(name='input',
+                           dtype='float64',
+                           shape=self.input_shape)
+            index = fluid.data(name='index',
+                               dtype='int64',
+                               shape=self.index_shape)
+            self.assertRaises(ValueError, paddle.index_select, x, index)
+
+    def test_dygraph_index_error(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.input_np)
+        index = paddle.to_tensor(self.index_np, dtype='int64')
+        self.assertRaises(ValueError, paddle.index_select, x, index)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 9dc291baf0520..5a608bc1874b8 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4658,13 +4658,13 @@ def frac(x, name=None):
 
 def take(input, index, name=None):
     """
-    Returns a new tensor with the elements of input at the given indices. 
+    Returns a new tensor with the elements of input at the given index.
     The input tensor is treated as if it were viewed as a 1-D tensor. 
-    The result takes the same shape as the indices.
+    The result takes the same shape as the index.
     
     Args:
-        input (Tensor): An N-D Tensor, which data type should be int32, int64, float32, float64. 
-        index (Tensor): An N-D Tensor, which data type should be int32, int64.
+        input (Tensor): An N-D Tensor, its data type should be int32, int64, float32, float64.
+        index (Tensor): An N-D Tensor, its data type should be int32, int64.
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4675,7 +4675,8 @@ def take(input, index, name=None):
 
             import paddle
 
-            x = paddle.arange(0, 12).reshape([3, 4])
+            x1 = paddle.arange(0, 12).reshape([3, 4])
+            x2 = paddle.arange(0, 12).reshape([3, 4])
             idx = paddle.arange(4, 10).reshape([2, 3])
 
             paddle.take(x, idx)

From 10b41c431fe93abd77fdfb70f5ac8c848b99af94 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Fri, 5 Aug 2022 08:54:00 +0800
Subject: [PATCH 05/16] fix Examples

---
 python/paddle/tensor/math.py | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index f401a51b1407e..6ade51f693e31 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4718,18 +4718,32 @@ def take(input, index, name=None):
     Examples:
         .. code-block:: python
 
+            import numpy as np
             import paddle
 
-            x1 = paddle.arange(0, 12).reshape([3, 4])
-            x2 = paddle.arange(0, 12).reshape([3, 4])
-            idx = paddle.arange(4, 10).reshape([2, 3])
+            n = np.arange(0, 12).reshape([3, 4])
+            x_int = paddle.to_tensor(n, dtype='int64')
+            x_float = paddle.to_tensor(n, dtype='float64')
+
+            idx_pos = paddle.arange(4, 10).reshape([2, 3])  # positive index
+            idx_neg = paddle.arange(-2, 4).reshape([2, 3])  # negative index
 
-            paddle.take(x, idx)
+            paddle.take(x_int, idx_pos)
             # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
             #        [[4, 5, 6],
             #         [7, 8, 9]])
 
-            x.take(idx)
+            paddle.take(x_int, idx_neg)
+            # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
+            #        [[10, 11, 0 ],
+            #         [1 , 2 , 3 ]])
+
+            paddle.take(x_float, idx_pos)
+            # Tensor(shape=[2, 3], dtype=float64, place=Place(cpu), stop_gradient=True,
+            #        [[4., 5., 6.],
+            #         [7., 8., 9.]])
+
+            x_int.take(idx_pos)
             # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
             #        [[4, 5, 6],
             #         [7, 8, 9]])

From 9649b876a1be7a22ad11e3ef000e36a9b5267dec Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Fri, 5 Aug 2022 16:29:47 +0800
Subject: [PATCH 06/16] fix Examples

---
 python/paddle/tensor/math.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 6ade51f693e31..42b7495c2ffd2 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4718,12 +4718,10 @@ def take(input, index, name=None):
     Examples:
         .. code-block:: python
 
-            import numpy as np
             import paddle
 
-            n = np.arange(0, 12).reshape([3, 4])
-            x_int = paddle.to_tensor(n, dtype='int64')
-            x_float = paddle.to_tensor(n, dtype='float64')
+            x_int = paddle.arange(0, 12).reshape([3, 4])
+            x_float = x_int.astype(paddle.float64)
 
             idx_pos = paddle.arange(4, 10).reshape([2, 3])  # positive index
             idx_neg = paddle.arange(-2, 4).reshape([2, 3])  # negative index

From 6806a8f579fee48d493a3799bf23ed244f94a207 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Mon, 22 Aug 2022 21:44:51 +0800
Subject: [PATCH 07/16] add param mode to take api

---
 .../paddle/fluid/tests/unittests/test_take.py | 95 ++++++++++++++-----
 python/paddle/tensor/math.py                  | 40 ++++++--
 2 files changed, 105 insertions(+), 30 deletions(-)

diff --git a/python/paddle/fluid/tests/unittests/test_take.py b/python/paddle/fluid/tests/unittests/test_take.py
index a003f72def9eb..0c2c097d20390 100644
--- a/python/paddle/fluid/tests/unittests/test_take.py
+++ b/python/paddle/fluid/tests/unittests/test_take.py
@@ -24,14 +24,14 @@
 
 class TestTakeAPI(unittest.TestCase):
 
+    def set_mode(self):
+        self.mode = 'raise'
+
     def set_dtype(self):
         self.input_dtype = 'float64'
         self.index_dtype = 'int64'
 
-    def setUp(self):
-        self.set_dtype()
-        self.place = fluid.CUDAPlace(
-            0) if core.is_compiled_with_cuda() else fluid.CPUPlace()
+    def set_input(self):
         self.input_shape = [3, 4]
         self.index_shape = [2, 3]
         self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
@@ -39,6 +39,13 @@ def setUp(self):
         self.index_np = np.arange(-4, 2).reshape(self.index_shape).astype(
             self.index_dtype)
 
+    def setUp(self):
+        self.set_mode()
+        self.set_dtype()
+        self.set_input()
+        self.place = fluid.CUDAPlace(
+            0) if core.is_compiled_with_cuda() else fluid.CPUPlace()
+
     def test_static_graph(self):
         paddle.enable_static()
         startup_program = Program()
@@ -50,7 +57,7 @@ def test_static_graph(self):
             index = fluid.data(name='index',
                                dtype=self.index_dtype,
                                shape=self.index_shape)
-            out = paddle.take(x, index)
+            out = paddle.take(x, index, mode=self.mode)
 
             exe = fluid.Executor(self.place)
             st_result = exe.run(fluid.default_main_program(),
@@ -60,15 +67,17 @@ def test_static_graph(self):
                                 },
                                 fetch_list=[out])
             self.assertTrue(
-                np.allclose(st_result, np.take(self.input_np, self.index_np)))
+                np.allclose(
+                    st_result,
+                    np.take(self.input_np, self.index_np, mode=self.mode)))
 
     def test_dygraph(self):
         paddle.disable_static(self.place)
         x = paddle.to_tensor(self.input_np)
         index = paddle.to_tensor(self.index_np)
-        dy_result = paddle.take(x, index)
+        dy_result = paddle.take(x, index, mode=self.mode)
         self.assertTrue(
-            np.allclose(np.take(self.input_np, self.index_np),
+            np.allclose(np.take(self.input_np, self.index_np, mode=self.mode),
                         dy_result.numpy()))
 
 
@@ -106,12 +115,13 @@ def test_static_type_error(self):
             x = fluid.data(name='input',
                            dtype=self.input_dtype,
                            shape=self.input_shape)
-            self.assertRaises(TypeError, paddle.take, x, self.index_np)
+            self.assertRaises(TypeError, paddle.take, x, self.index_np,
+                              self.mode)
 
     def test_dygraph_type_error(self):
         paddle.disable_static(self.place)
         x = paddle.to_tensor(self.input_np)
-        self.assertRaises(TypeError, paddle.take, x, self.index_np)
+        self.assertRaises(TypeError, paddle.take, x, self.index_np, self.mode)
 
     def test_static_dtype_error(self):
         """Data type of argument 'index' must be in [paddle.int32, paddle.int64]"""
@@ -123,28 +133,39 @@ def test_static_dtype_error(self):
             index = fluid.data(name='index',
                                dtype='float32',
                                shape=self.index_shape)
-            self.assertRaises(TypeError, paddle.take, x, index)
+            self.assertRaises(TypeError, paddle.take, x, index, self.mode)
 
     def test_dygraph_dtype_error(self):
         paddle.disable_static(self.place)
         x = paddle.to_tensor(self.input_np)
         index = paddle.to_tensor(self.index_np, dtype='float32')
-        self.assertRaises(TypeError, paddle.take, x, index)
+        self.assertRaises(TypeError, paddle.take, x, index, self.mode)
 
 
-class TestTakeIndexRangeError(TestTakeAPI):
+class TestTakeModeRaise(unittest.TestCase):
     """Test take index out of range error"""
 
+    def set_mode(self):
+        self.mode = 'raise'
+
+    def set_dtype(self):
+        self.input_dtype = 'float64'
+        self.index_dtype = 'int64'
+
+    def set_input(self):
+        self.input_shape = [3, 4]
+        self.index_shape = [5, 8]
+        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
+            self.input_dtype)
+        self.index_np = np.arange(-20, 20).reshape(self.index_shape).astype(
+            self.index_dtype)  # Both ends of the index are out of bounds
+
     def setUp(self):
+        self.set_mode()
         self.set_dtype()
+        self.set_input()
         self.place = fluid.CUDAPlace(
             0) if core.is_compiled_with_cuda() else fluid.CPUPlace()
-        self.input_shape = [3, 4]
-        self.index_shape = [2, 3]
-        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
-            self.input_dtype)
-        self.index_np = np.arange(6, 12).reshape(self.index_shape).astype(
-            self.index_dtype)
 
     def test_static_index_error(self):
         """When the index is out of range,
@@ -152,19 +173,49 @@ def test_static_index_error(self):
         paddle.enable_static()
         with program_guard(Program()):
             x = fluid.data(name='input',
-                           dtype='float64',
+                           dtype=self.input_dtype,
                            shape=self.input_shape)
             index = fluid.data(name='index',
-                               dtype='int64',
+                               dtype=self.index_dtype,
                                shape=self.index_shape)
             self.assertRaises(ValueError, paddle.index_select, x, index)
 
     def test_dygraph_index_error(self):
         paddle.disable_static(self.place)
         x = paddle.to_tensor(self.input_np)
-        index = paddle.to_tensor(self.index_np, dtype='int64')
+        index = paddle.to_tensor(self.index_np, dtype=self.index_dtype)
         self.assertRaises(ValueError, paddle.index_select, x, index)
 
 
+class TestTakeModeWrap(TestTakeAPI):
+    """Test take index out of range mode"""
+
+    def set_mode(self):
+        self.mode = 'wrap'
+
+    def set_input(self):
+        self.input_shape = [3, 4]
+        self.index_shape = [5, 8]
+        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
+            self.input_dtype)
+        self.index_np = np.arange(-20, 20).reshape(self.index_shape).astype(
+            self.index_dtype)  # Both ends of the index are out of bounds
+
+
+class TestTakeModeClip(TestTakeAPI):
+    """Test take index out of range mode"""
+
+    def set_mode(self):
+        self.mode = 'clip'
+
+    def set_input(self):
+        self.input_shape = [3, 4]
+        self.index_shape = [5, 8]
+        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
+            self.input_dtype)
+        self.index_np = np.arange(-20, 20).reshape(self.index_shape).astype(
+            self.index_dtype)  # Both ends of the index are out of bounds
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 42b7495c2ffd2..50710321876fd 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4701,15 +4701,22 @@ def frac(x, name=None):
                 type="trunc", inputs=inputs, attrs=attrs, outputs={"Out": y})
             return _elementwise_op(LayerHelper(op_type, **locals()))
 
-def take(input, index, name=None):
+def take(x, index, mode='raise', name=None):
     """
-    Returns a new tensor with the elements of input at the given index.
-    The input tensor is treated as if it were viewed as a 1-D tensor. 
+    Returns a new tensor with the elements of tnput tensor x at the given index.
+    The input tensor is treated as if it were viewed as a 1-D tensor.
     The result takes the same shape as the index.
-    
+
     Args:
-        input (Tensor): An N-D Tensor, its data type should be int32, int64, float32, float64.
+        x (Tensor): An N-D Tensor, its data type should be int32, int64, float32, float64.
         index (Tensor): An N-D Tensor, its data type should be int32, int64.
+        mode (str, optional): Specifies how out-of-bounds index will behave.
+                the candicates are ``'raise'`` | ``'wrap'`` | ``'clip'``.
+                If :attr:`mode` is ``'raise'``, raise an error (default);
+                If :attr:`mode` is ``'wrap'``, wrap around;
+                If :attr:`mode` is ``'clip'``, clip to the range.
+                    ``'clip'`` mode means that all indices that are too large are replaced by the index that
+                    addresses the last element. Note that this disables indexing with negative numbers.
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4746,6 +4753,9 @@ def take(input, index, name=None):
             #        [[4, 5, 6],
             #         [7, 8, 9]])
     """
+    if mode not in ['raise', 'wrap', 'clip']:
+        raise ValueError(
+            "'mode' in 'take' should be 'raise', 'wrap', 'clip', but received {}.".format(mode))
 
     if paddle.in_dynamic_mode():
         if not isinstance(index, (paddle.Tensor, Variable)):
@@ -4755,14 +4765,28 @@ def take(input, index, name=None):
             raise TypeError(
                 "The data type of 'index' must be one of ['int32', 'int64'], but got {}".format(
                     index.dtype))
+
     else:
         check_variable_and_dtype(index, 'index', ['int32', 'int64'], 'take')
 
-    input_1d = input.flatten()
+    input_1d = x.flatten()
     index_1d = index.flatten()
+    max_index = input_1d.shape[-1]
+
+    if mode == 'raise':
+        # This processing enables 'take' to handle negative indexes within the correct range.
+        index_1d = paddle.where(index_1d < 0, index_1d % max_index, index_1d)
+        pass
+    elif mode == 'wrap':
+        # The out of range indices are constrained by taking the remainder.
+        index_1d = paddle.where(index_1d < 0,
+                                index_1d % max_index, index_1d)
+        index_1d = paddle.where(index_1d >= max_index,
+                                index_1d % max_index, index_1d)
+    elif mode == 'clip':
+        # 'clip' mode disables indexing with negative numbers.
+        index_1d = clip(index_1d, 0, max_index - 1)
 
-    # This processing enables 'take' to handle negative indexes within the correct range
-    index_1d = paddle.where(index_1d < 0, index_1d + input_1d.shape[0], index_1d)
     out = input_1d.index_select(index_1d).reshape(index.shape)
 
     return out

From 5d32c521c2862d618988370d0fb6772f171c4f49 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Mon, 22 Aug 2022 22:00:52 +0800
Subject: [PATCH 08/16] add example code

---
 python/paddle/tensor/math.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 71a245ab3b23c..e98c7b0e984db 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4828,6 +4828,7 @@ def take(x, index, mode='raise', name=None):
 
             idx_pos = paddle.arange(4, 10).reshape([2, 3])  # positive index
             idx_neg = paddle.arange(-2, 4).reshape([2, 3])  # negative index
+            idx_err = paddle.arange(-2, 13).reshape([3, 5])  # index out of range
 
             paddle.take(x_int, idx_pos)
             # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
@@ -4848,6 +4849,19 @@ def take(x, index, mode='raise', name=None):
             # Tensor(shape=[2, 3], dtype=int64, place=Place(cpu), stop_gradient=True,
             #        [[4, 5, 6],
             #         [7, 8, 9]])
+
+            paddle.take(x_int, idx_err, mode='wrap')
+            # Tensor(shape=[3, 5], dtype=int32, place=Place(cpu), stop_gradient=True,
+            #        [[10, 11, 0 , 1 , 2 ],
+            #         [3 , 4 , 5 , 6 , 7 ],
+            #         [8 , 9 , 10, 11, 0 ]])
+
+            paddle.take(x_int, idx_err, mode='clip')
+            # Tensor(shape=[3, 5], dtype=int32, place=Place(cpu), stop_gradient=True,
+            #        [[0 , 0 , 0 , 1 , 2 ],
+            #         [3 , 4 , 5 , 6 , 7 ],
+            #         [8 , 9 , 10, 11, 11]])
+
     """
     if mode not in ['raise', 'wrap', 'clip']:
         raise ValueError(

From b35d831bc423e59e825353ffe941584ddeee89b0 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Tue, 23 Aug 2022 15:09:28 +0800
Subject: [PATCH 09/16] fix test using np.testing.assert_allclose

---
 python/paddle/fluid/tests/unittests/test_take.py | 15 +++++++--------
 python/paddle/tensor/math.py                     |  1 -
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/python/paddle/fluid/tests/unittests/test_take.py b/python/paddle/fluid/tests/unittests/test_take.py
index 0c2c097d20390..ba04d1636ba5b 100644
--- a/python/paddle/fluid/tests/unittests/test_take.py
+++ b/python/paddle/fluid/tests/unittests/test_take.py
@@ -65,20 +65,19 @@ def test_static_graph(self):
                                     'input': self.input_np,
                                     'index': self.index_np
                                 },
-                                fetch_list=[out])
-            self.assertTrue(
-                np.allclose(
-                    st_result,
-                    np.take(self.input_np, self.index_np, mode=self.mode)))
+                                fetch_list=out)
+            np.testing.assert_allclose(
+                st_result[0],
+                np.take(self.input_np, self.index_np, mode=self.mode))
 
     def test_dygraph(self):
         paddle.disable_static(self.place)
         x = paddle.to_tensor(self.input_np)
         index = paddle.to_tensor(self.index_np)
         dy_result = paddle.take(x, index, mode=self.mode)
-        self.assertTrue(
-            np.allclose(np.take(self.input_np, self.index_np, mode=self.mode),
-                        dy_result.numpy()))
+        np.testing.assert_allclose(
+            np.take(self.input_np, self.index_np, mode=self.mode),
+            dy_result.numpy())
 
 
 class TestTakeInt32(TestTakeAPI):
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index e98c7b0e984db..c8ebf67c4c981 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4886,7 +4886,6 @@ def take(x, index, mode='raise', name=None):
     if mode == 'raise':
         # This processing enables 'take' to handle negative indexes within the correct range.
         index_1d = paddle.where(index_1d < 0, index_1d % max_index, index_1d)
-        pass
     elif mode == 'wrap':
         # The out of range indices are constrained by taking the remainder.
         index_1d = paddle.where(index_1d < 0,

From c4161f29cbd9b98a0db72442218a64bfb937bf2d Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Tue, 23 Aug 2022 17:11:40 +0800
Subject: [PATCH 10/16] add annotation

---
 python/paddle/tensor/math.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index c8ebf67c4c981..bb1fcfcb093b6 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4885,6 +4885,8 @@ def take(x, index, mode='raise', name=None):
 
     if mode == 'raise':
         # This processing enables 'take' to handle negative indexes within the correct range.
+        # Negative indexes can be enabled,
+        # but out-of-range indexes will report an error in the following paddle.index_select
         index_1d = paddle.where(index_1d < 0, index_1d % max_index, index_1d)
     elif mode == 'wrap':
         # The out of range indices are constrained by taking the remainder.

From ca2604f218e29a1213b430be7d5a3b760683a0dd Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Tue, 23 Aug 2022 17:54:45 +0800
Subject: [PATCH 11/16] fix typo

---
 python/paddle/tensor/math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index bb1fcfcb093b6..d04a1f678c02b 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4799,7 +4799,7 @@ def sgn(x, name=None):
 
 def take(x, index, mode='raise', name=None):
     """
-    Returns a new tensor with the elements of tnput tensor x at the given index.
+    Returns a new tensor with the elements of input tensor x at the given index.
     The input tensor is treated as if it were viewed as a 1-D tensor.
     The result takes the same shape as the index.
 

From 7b3fc1d1b1f3bbb8eea6dfde44813aa6f3b269a7 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Wed, 24 Aug 2022 15:14:30 +0800
Subject: [PATCH 12/16] =?UTF-8?q?fix=20=E5=B5=8C=E5=A5=97=E5=88=97?=
 =?UTF-8?q?=E8=A1=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 python/paddle/tensor/math.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index d04a1f678c02b..d1bccb896b19d 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4806,13 +4806,13 @@ def take(x, index, mode='raise', name=None):
     Args:
         x (Tensor): An N-D Tensor, its data type should be int32, int64, float32, float64.
         index (Tensor): An N-D Tensor, its data type should be int32, int64.
-        mode (str, optional): Specifies how out-of-bounds index will behave.
-                the candicates are ``'raise'`` | ``'wrap'`` | ``'clip'``.
-                If :attr:`mode` is ``'raise'``, raise an error (default);
-                If :attr:`mode` is ``'wrap'``, wrap around;
-                If :attr:`mode` is ``'clip'``, clip to the range.
-                    ``'clip'`` mode means that all indices that are too large are replaced by the index that
-                    addresses the last element. Note that this disables indexing with negative numbers.
+        mode (str, optional): Specifies how out-of-bounds index will behave. the candicates are ``'raise'``, ``'wrap'`` and ``'clip'``.
+
+            - ``'raise'``: raise an error (default);
+            - ``'wrap'``: wrap around;
+            - ``'clip'``: clip to the range. ``'clip'`` mode means that all indices that are too large are replaced by
+            the index that addresses the last element. Note that this disables indexing with negative numbers.
+
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:

From 64b688a53aefd2f6dafc9838601d41bb719cfd43 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Wed, 24 Aug 2022 15:32:24 +0800
Subject: [PATCH 13/16] fix Tensor,

---
 python/paddle/tensor/math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 457deccebdafc..9b0ee8c675428 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4795,7 +4795,7 @@ def take(x, index, mode='raise', name=None):
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
-        Tensor: Tensor with the same shape as index, the data type is the same with input.
+        Tensor, Tensor with the same shape as index, the data type is the same with input.
 
     Examples:
         .. code-block:: python

From cdd1080b683099ee3ea841ca4528e6f48d18353c Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Thu, 25 Aug 2022 20:35:07 +0800
Subject: [PATCH 14/16] fix docs warning

---
 python/paddle/tensor/math.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 9b0ee8c675428..535d51ea668a8 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4789,8 +4789,7 @@ def take(x, index, mode='raise', name=None):
 
             - ``'raise'``: raise an error (default);
             - ``'wrap'``: wrap around;
-            - ``'clip'``: clip to the range. ``'clip'`` mode means that all indices that are too large are replaced by
-            the index that addresses the last element. Note that this disables indexing with negative numbers.
+            - ``'clip'``: clip to the range. ``'clip'`` mode means that all indices that are too large are replaced by the index that addresses the last element. Note that this disables indexing with negative numbers.
 
         name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 

From 4ca5c417c69ab0042702ecd13fd5d128571d9717 Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Sat, 27 Aug 2022 12:38:36 +0800
Subject: [PATCH 15/16] fix raise bug

---
 python/paddle/tensor/math.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
index 1d15475572f65..9dd1ab36fd0f3 100644
--- a/python/paddle/tensor/math.py
+++ b/python/paddle/tensor/math.py
@@ -4863,9 +4863,7 @@ def take(x, index, mode='raise', name=None):
 
     if mode == 'raise':
         # This processing enables 'take' to handle negative indexes within the correct range.
-        # Negative indexes can be enabled,
-        # but out-of-range indexes will report an error in the following paddle.index_select
-        index_1d = paddle.where(index_1d < 0, index_1d % max_index, index_1d)
+        index_1d = paddle.where(index_1d < 0, index_1d + max_index, index_1d)
     elif mode == 'wrap':
         # The out of range indices are constrained by taking the remainder.
         index_1d = paddle.where(index_1d < 0,

From 7fd6c853883cc48ff037509c1c28ad87a4a93a2e Mon Sep 17 00:00:00 2001
From: S-HuaBomb <s_huabang@163.com>
Date: Mon, 29 Aug 2022 11:13:17 +0800
Subject: [PATCH 16/16] add test case for negative index out of range error

---
 .../paddle/fluid/tests/unittests/test_take.py | 36 ++++++++++++++++---
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/python/paddle/fluid/tests/unittests/test_take.py b/python/paddle/fluid/tests/unittests/test_take.py
index ba04d1636ba5b..6e58a3a43de47 100644
--- a/python/paddle/fluid/tests/unittests/test_take.py
+++ b/python/paddle/fluid/tests/unittests/test_take.py
@@ -141,8 +141,8 @@ def test_dygraph_dtype_error(self):
         self.assertRaises(TypeError, paddle.take, x, index, self.mode)
 
 
-class TestTakeModeRaise(unittest.TestCase):
-    """Test take index out of range error"""
+class TestTakeModeRaisePos(unittest.TestCase):
+    """Test positive index out of range error"""
 
     def set_mode(self):
         self.mode = 'raise'
@@ -153,11 +153,11 @@ def set_dtype(self):
 
     def set_input(self):
         self.input_shape = [3, 4]
-        self.index_shape = [5, 8]
+        self.index_shape = [5, 6]
         self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
             self.input_dtype)
-        self.index_np = np.arange(-20, 20).reshape(self.index_shape).astype(
-            self.index_dtype)  # Both ends of the index are out of bounds
+        self.index_np = np.arange(-10, 20).reshape(self.index_shape).astype(
+            self.index_dtype)  # positive indices are out of range
 
     def setUp(self):
         self.set_mode()
@@ -186,6 +186,32 @@ def test_dygraph_index_error(self):
         self.assertRaises(ValueError, paddle.index_select, x, index)
 
 
+class TestTakeModeRaiseNeg(TestTakeModeRaisePos):
+    """Test negative index out of range error"""
+
+    def set_mode(self):
+        self.mode = 'raise'
+
+    def set_dtype(self):
+        self.input_dtype = 'float64'
+        self.index_dtype = 'int64'
+
+    def set_input(self):
+        self.input_shape = [3, 4]
+        self.index_shape = [5, 6]
+        self.input_np = np.arange(0, 12).reshape(self.input_shape).astype(
+            self.input_dtype)
+        self.index_np = np.arange(-20, 10).reshape(self.index_shape).astype(
+            self.index_dtype)  # negative indices are out of range
+
+    def setUp(self):
+        self.set_mode()
+        self.set_dtype()
+        self.set_input()
+        self.place = fluid.CUDAPlace(
+            0) if core.is_compiled_with_cuda() else fluid.CPUPlace()
+
+
 class TestTakeModeWrap(TestTakeAPI):
     """Test take index out of range mode"""