diff --git a/python/paddle/fluid/tests/unittests/test_index_select_op.py b/python/paddle/fluid/tests/unittests/test_index_select_op.py index 50d04c1a72378d4491b3570fd687a76f045b9551..e551989ed322db0a49e1e95338aeb5fb356a4951 100644 --- a/python/paddle/fluid/tests/unittests/test_index_select_op.py +++ b/python/paddle/fluid/tests/unittests/test_index_select_op.py @@ -83,7 +83,7 @@ class TestIndexSelectAPI(unittest.TestCase): x = fluid.layers.data(name='x', shape=[-1, 4]) index = fluid.layers.data( name='index', shape=[3], dtype='int32', append_batch_size=False) - z = paddle.index_select(x, index, dim=1) + z = paddle.index_select(x, index, axis=1) exe = fluid.Executor(fluid.CPUPlace()) res, = exe.run(feed={'x': self.data_x, 'index': self.data_index}, @@ -124,7 +124,7 @@ class TestIndexSelectAPI(unittest.TestCase): with fluid.dygraph.guard(): x = fluid.dygraph.to_variable(self.data_x) index = fluid.dygraph.to_variable(self.data_index) - z = paddle.index_select(x, index, dim=1) + z = paddle.index_select(x, index, axis=1) np_z = z.numpy() expect_out = np.array([[1.0, 2.0, 2.0], [5.0, 6.0, 6.0], [9.0, 10.0, 10.0]]) diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index 59b8f1e765b266a2fda61dd645b9fdce96e2a400..d8874e47c357937020ffe2e392332599df6653c0 100644 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -63,9 +63,9 @@ def argmax(input, axis=None, dtype=None, out=None, keepdims=False, name=None): Variable that meets the requirements to store the result of operation. if out is None, a new Varibale will be create to store the result. Defalut is None. keepdims(bool, optional): Keep the axis that do the select max. - name(str, optional): The name of output variable, normally there is no need for user to set this this property. - Default value is None, the framework set the name of output variable. - + name(str, optional): The default value is None. Normally there is no + need for user to set this property. For more information, please + refer to :ref:`api_guide_Name`. Returns: Variable: A Tensor with data type int64. @@ -135,7 +135,7 @@ def argmax(input, axis=None, dtype=None, out=None, keepdims=False, name=None): return out -def index_select(input, index, dim=0): +def index_select(x, index, axis=0, name=None): """ :alias_main: paddle.index_select :alias: paddle.index_select,paddle.tensor.index_select,paddle.tensor.search.index_select @@ -146,56 +146,60 @@ def index_select(input, index, dim=0): size as the length of `index`; other dimensions have the same size as in the `input` tensor. Args: - input (Variable): The input tensor variable. - index (Variable): The 1-D tensor containing the indices to index. - dim (int): The dimension in which we index. + x (Variable): The input tensor variable.The dtype of x can be one of float32, float64, int32, int64. + index (Variable): The 1-D tensor containing the indices to index.the dtype of index can be int32 or int64. + axis (int, optional): The dimension in which we index. Default: if None, the axis is 0. + name(str, optional): The default value is None. Normally there is no + need for user to set this property. For more information, please + refer to :ref:`api_guide_Name`. Returns: Variable: A Tensor with same data type as `input`. + + Raises: + TypeError: x must be a Variable and the dtype of x must be one of float32, float64, int32 and int64. + TypeError: index must be a Variable adn the dtype of index must be int32 or int64. Examples: .. code-block:: python import paddle - import paddle.fluid as fluid import numpy as np + paddle.enable_imperative() # Now we are in imperative mode data = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]) data_index = np.array([0, 1, 1]).astype('int32') - with fluid.dygraph.guard(): - x = fluid.dygraph.to_variable(data) - index = fluid.dygraph.to_variable(data_index) - out_z1 = paddle.index_select(x, index) - print(out_z1.numpy()) - #[[1. 2. 3. 4.] - # [5. 6. 7. 8.] - # [5. 6. 7. 8.]] - out_z2 = paddle.index_select(x, index, dim=1) - print(out_z2.numpy()) - #[[ 1. 2. 2.] - # [ 5. 6. 6.] - # [ 9. 10. 10.]] + x = paddle.imperative.to_variable(data) + index = paddle.imperative.to_variable(data_index) + out_z1 = paddle.index_select(x=x, index=index) + #[[1. 2. 3. 4.] + # [5. 6. 7. 8.] + # [5. 6. 7. 8.]] + out_z2 = paddle.index_select(x=x, index=index, axis=1) + #[[ 1. 2. 2.] + # [ 5. 6. 6.] + # [ 9. 10. 10.]] """ - helper = LayerHelper("index_select", **locals()) + if in_dygraph_mode(): - return core.ops.index_select(input, index, 'dim', dim) + return core.ops.index_select(x, index, 'dim', axis) - check_variable_and_dtype(input, 'x', - ['float32', 'float64', 'int32', 'int64'], - 'paddle.tensor.search.index_sample') + helper = LayerHelper("index_select", **locals()) + check_variable_and_dtype(x, 'x', ['float32', 'float64', 'int32', 'int64'], + 'paddle.tensor.search.index_select') check_variable_and_dtype(index, 'index', ['int32', 'int64'], - 'paddle.tensor.search.index_sample') + 'paddle.tensor.search.index_select') - out = helper.create_variable_for_type_inference(input.dtype) + out = helper.create_variable_for_type_inference(x.dtype) helper.append_op( type='index_select', - inputs={'X': input, + inputs={'X': x, 'Index': index}, outputs={'Out': out}, - attrs={'dim': dim}) + attrs={'dim': axis}) return out