diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py index d2e8ae581a030b54ac9fea0eea24e325e1266125..8a0ec03986b5790f626d2746b84451ae89453a97 100644 --- a/python/paddle/__init__.py +++ b/python/paddle/__init__.py @@ -34,6 +34,7 @@ import paddle.distributed batch = batch.batch import paddle.sysconfig import paddle.nn +import paddle.tensor #TODO: define alias in tensor and framework directory # from .tensor.creation import create_.tensor #DEFINE_ALIAS @@ -191,7 +192,7 @@ from .tensor.manipulation import flip #DEFINE_ALIAS # from .tensor.search import where #DEFINE_ALIAS # from .tensor.search import index_select #DEFINE_ALIAS # from .tensor.search import nonzero #DEFINE_ALIAS -# from .tensor.search import sort #DEFINE_ALIAS +from .tensor.search import sort #DEFINE_ALIAS # from .framework.framework import set_default_dtype #DEFINE_ALIAS # from .framework.framework import get_default_dtype #DEFINE_ALIAS # from .framework.random import manual_seed #DEFINE_ALIAS diff --git a/python/paddle/fluid/tests/unittests/test_argsort_op.py b/python/paddle/fluid/tests/unittests/test_argsort_op.py index 140502e896c7d66b51d8a2ea769fc2164d69c387..758338743625eda4ea632aca590cbe3a7a43de36 100644 --- a/python/paddle/fluid/tests/unittests/test_argsort_op.py +++ b/python/paddle/fluid/tests/unittests/test_argsort_op.py @@ -15,6 +15,7 @@ from __future__ import print_function import unittest +import paddle import paddle.fluid as fluid import paddle.fluid.layers as layers import numpy as np @@ -320,5 +321,35 @@ class TestArgsortOpDescendingAxisNeg2GPU(TestArgsortOpAxisNeg2GPU): self.descending = True +class TestSortOnCPU(TestArgsortOpCPU): + def init_place(self): + self.place = core.CPUPlace() + + def test_out(self): + self.init_place() + with fluid.program_guard(fluid.Program()): + input = fluid.data(name="input", shape=[2, 3, 4], dtype="float32") + + res = fluid.data(name="output", shape=[2, 3, 4], dtype="float32") + output = paddle.tensor.sort(input=input, out=res) + + exe = fluid.Executor(self.place) + data = np.array( + [[[5, 8, 9, 5], [0, 0, 1, 7], [6, 9, 2, 4]], + [[5, 2, 4, 2], [4, 7, 7, 9], [1, 7, 0, 6]]], + dtype='float32') + result = exe.run(feed={'input': data}, fetch_list=[res, output[0]]) + + self.assertEqual((result[0] == result[1]).all(), True) + + +class TestSortOnGPU(TestSortOnCPU): + def init_place(self): + if core.is_compiled_with_cuda(): + self.place = core.CUDAPlace(0) + else: + self.place = core.CPUPlace() + + if __name__ == "__main__": unittest.main() diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py index 65758e0a9218237c8a1667c781b11aff281210cf..51cf5da6df18c04e40b43c14c22f0572bb63968f 100644 --- a/python/paddle/tensor/__init__.py +++ b/python/paddle/tensor/__init__.py @@ -168,4 +168,4 @@ from .manipulation import flip #DEFINE_ALIAS # from .search import where #DEFINE_ALIAS # from .search import index_select #DEFINE_ALIAS # from .search import nonzero #DEFINE_ALIAS -# from .search import sort #DEFINE_ALIAS +from .search import sort #DEFINE_ALIAS diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index bfc607908abe65226969b49d44f259bb34f3e35d..59c89797bb621c68fee81445788e0bf2b443d8c5 100644 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -13,14 +13,108 @@ # limitations under the License. # TODO: define searching & indexing functions of a tensor -# __all__ = ['argmax', -# 'argmin', -# 'argsort', -# 'has_inf', -# 'has_nan', -# 'masked_select', -# 'topk', -# 'where', -# 'index_select', -# 'nonzero', -# 'sort'] +__all__ = [ + 'argmax', + # 'argmin', + # 'argsort', + # 'has_inf', + # 'has_nan', + # 'masked_select', + # 'topk', + # 'where', + # 'index_select', + # 'nonzero', + 'sort' +] + +from paddle.common_ops_import import * +import warnings + + +def sort(input, axis=-1, descending=False, out=None, name=None): + """ + This OP sorts the input along the given axis, and returns sorted output + data Varibale and its corresponding index Variable with the same shape as + :attr:`input`. + + **NOTICE**: The Variable in the output of this OP has gradient. You could\ + set Variable :attr:`stop_gradient`. + Args: + input(Variable): An input N-D Tensor with type float32, float64, int16, + int32, int64, uint8. + axis(int, optional): Axis to compute indices along. The effective range + is [-R, R), where R is Rank(x). when axis<0, it works the same way + as axis+R. Default is 0. + descending(bool, optional) : Descending is a flag, if set to true, + algorithm will sort by descending order, else sort by + ascending order. Default is false. + out(Variable, optional): The default value is None. Optional output + which can be any created 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. + 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: + tuple: A tuple of sorted data Variable(with the same shape and data + type as input) and the sorted indices(with the same shape as input's + and with data type int64). + Examples: + .. code-block:: python + import paddle + import paddle.fluid as fluid + import numpy as np + in1 = np.array([[[5,8,9,5], + [0,0,1,7], + [6,9,2,4]], + [[5,2,4,2], + [4,7,7,9], + [1,7,0,6]]]).astype(np.float32) + with fluid.dygraph.guard(): + x = fluid.dygraph.to_variable(in1) + out1 = paddle.sort(input=x, axis=-1) + out2 = paddle.sort(input=x, axis=0) + out3 = paddle.sort(input=x, axis=1) + print(out1[0].numpy()) + # [[[5. 5. 8. 9.] + # [0. 0. 1. 7.] + # [2. 4. 6. 9.]] + # [[2. 2. 4. 5.] + # [4. 7. 7. 9.] + # [0. 1. 6. 7.]]] + print(out1[1].numpy()) + # [[[0 3 1 2] + # [0 1 2 3] + # [2 3 0 1]] + # [[1 3 2 0] + # [0 1 2 3] + # [2 0 3 1]]] + print(out2[0].numpy()) + # [[[5. 2. 4. 2.] + # [0. 0. 1. 7.] + # [1. 7. 0. 4.]] + # [[5. 8. 9. 5.] + # [4. 7. 7. 9.] + # [6. 9. 2. 6.]]] + print(out3[0].numpy()) + # [[[0. 0. 1. 4.] + # [5. 8. 2. 5.] + # [6. 9. 9. 7.]] + # [[1. 2. 0. 2.] + # [4. 7. 4. 6.] + # [5. 7. 7. 9.]]] + """ + helper = LayerHelper("sort", **locals()) + if out is None: + out = helper.create_variable_for_type_inference( + dtype=input.dtype, stop_gradient=False) + ids = helper.create_variable_for_type_inference( + VarDesc.VarType.INT64, stop_gradient=True) + helper.append_op( + type='argsort', + inputs={'X': input}, + outputs={'Out': out, + 'Indices': ids}, + attrs={'axis': axis, + 'descending': descending}) + return out, ids