search.py 27.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
#   Copyright (c) 2020 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.
C
Chengmo 已提交
14
from __future__ import print_function
15
import numpy as np
C
Chengmo 已提交
16 17
from ..fluid.layer_helper import LayerHelper
from ..fluid.data_feeder import check_variable_and_dtype, check_type, check_dtype
18
from ..fluid import core, layers
19 20 21 22
from paddle.common_ops_import import in_dygraph_mode
from paddle.common_ops_import import convert_np_dtype_to_dtype_
from paddle.common_ops_import import Variable
from paddle.common_ops_import import VarDesc
23

24
# TODO: define searching & indexing functions of a tensor  
25 26
# from ..fluid.layers import has_inf  #DEFINE_ALIAS
# from ..fluid.layers import has_nan  #DEFINE_ALIAS
27

28 29
__all__ = []

30

31 32
def argsort(x, axis=-1, descending=False, name=None):
    """
W
wawltor 已提交
33
    This OP sorts the input along the given axis, and returns the corresponding index tensor for the sorted output values. The default sort algorithm is ascending, if you want the sort algorithm to be descending, you must set the :attr:`descending` as True.
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52

    Args:
        x(Tensor): 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.
        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:
        Tensor: sorted indices(with the same shape as ``x``
        and with data type int64).

    Examples:
李灿 已提交
53

54
        .. code-block:: python
李灿 已提交
55

56 57
            import paddle
            
58 59 60 61 62 63 64
            x = paddle.to_tensor([[[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')
65 66 67
            out1 = paddle.argsort(x=x, axis=-1)
            out2 = paddle.argsort(x=x, axis=0)
            out3 = paddle.argsort(x=x, axis=1)
N
Noel 已提交
68
            print(out1)
W
wawltor 已提交
69 70 71
            #[[[0 3 1 2]
            #  [0 1 2 3]
            #  [2 3 0 1]]
72
            # [[1 3 2 0]
W
wawltor 已提交
73 74
            #  [0 1 2 3]
            #  [2 0 3 1]]]
N
Noel 已提交
75
            print(out2)
W
wawltor 已提交
76 77 78 79 80 81
            #[[[0 1 1 1]
            #  [0 0 0 0]
            #  [1 1 1 0]]
            # [[1 0 0 0]
            #  [1 1 1 1]
            #  [0 0 0 1]]]
N
Noel 已提交
82
            print(out3)
W
wawltor 已提交
83 84 85 86 87 88
            #[[[1 1 1 2]
            #  [0 0 2 0]
            #  [2 2 0 1]]
            # [[2 0 2 0]
            #  [1 1 0 2]
            #  [0 2 1 1]]]
89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
    """
    if in_dygraph_mode():
        _, ids = core.ops.argsort(x, 'axis', axis, 'descending', descending)
        return ids
    check_variable_and_dtype(
        x, 'x', ['float32', 'float64', 'int16', 'int32', 'int64', 'uint8'],
        'argsort')

    helper = LayerHelper("argsort", **locals())
    out = helper.create_variable_for_type_inference(
        dtype=x.dtype, stop_gradient=True)
    ids = helper.create_variable_for_type_inference(
        VarDesc.VarType.INT64, stop_gradient=True)
    helper.append_op(
        type='argsort',
        inputs={'X': x},
        outputs={'Out': out,
                 'Indices': ids},
        attrs={'axis': axis,
               'descending': descending})
    return ids


112
def argmax(x, axis=None, keepdim=False, dtype="int64", name=None):
113 114 115 116 117
    """
    This OP computes the indices of the max elements of the input tensor's
    element along the provided axis.

    Args:
W
wawltor 已提交
118
        x(Tensor): An input N-D Tensor with type float32, float64, int16,
119 120
            int32, int64, uint8.
        axis(int, optional): Axis to compute indices along. The effective range
W
wawltor 已提交
121 122 123
            is [-R, R), where R is x.ndim. when axis < 0, it works the same way
            as axis + R. Default is None, the input `x` will be into the flatten tensor, and selecting the min value index.
        keepdim(bool, optional): Keep the axis that selecting max. The defalut value is False.
124 125 126
        dtype(str|np.dtype, optional): Data type of the output tensor which can
                    be int32, int64. The default value is 'int64', and it will
                    return the int64 indices.
127 128 129
        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`.
130 131

    Returns:
W
wawltor 已提交
132
        Tensor, return the tensor of `int32` if set :attr:`dtype` is `int32`, otherwise return the tensor of `int64`
133 134 135 136

    Examples:
        .. code-block:: python

W
wawltor 已提交
137
            import paddle
138

139 140 141
            x =  paddle.to_tensor([[5,8,9,5],
                                     [0,0,1,7],
                                     [6,9,2,4]])
W
wawltor 已提交
142
            out1 = paddle.argmax(x)
N
Noel 已提交
143
            print(out1) # 2
W
wawltor 已提交
144
            out2 = paddle.argmax(x, axis=1)
N
Noel 已提交
145
            print(out2) 
W
wawltor 已提交
146 147
            # [2 3 1]
            out3 = paddle.argmax(x, axis=-1)
N
Noel 已提交
148
            print(out3) 
W
wawltor 已提交
149
            # [2 3 1]
150
    """
151 152 153 154
    if axis is not None and not isinstance(axis, int):
        raise TypeError(
            "The type of 'axis'  must be int or None in argmax, but received %s."
            % (type(axis)))
155

156 157 158 159
    if dtype is None:
        raise ValueError(
            "the value of 'dtype' in argmax could not be None, but received None"
        )
160

161 162
    var_dtype = convert_np_dtype_to_dtype_(dtype)
    check_dtype(var_dtype, 'dtype', ['int32', 'int64'], 'argmin')
W
wawltor 已提交
163 164 165 166 167 168
    flatten = False
    if axis is None:
        flatten = True
        axis = 0

    if in_dygraph_mode():
169 170
        out = core.ops.arg_max(x, 'axis', axis, 'dtype', var_dtype, 'keepdims',
                               keepdim, 'flatten', flatten)
W
wawltor 已提交
171 172 173 174 175 176
        return out

    helper = LayerHelper("argmax", **locals())
    check_variable_and_dtype(
        x, 'x', ['float32', 'float64', 'int16', 'int32', 'int64', 'uint8'],
        'paddle.argmax')
177
    attrs = {}
W
wawltor 已提交
178 179 180 181
    out = helper.create_variable_for_type_inference(var_dtype)
    attrs['keepdims'] = keepdim
    attrs['axis'] = axis
    attrs['flatten'] = flatten
182
    attrs['dtype'] = var_dtype
W
wawltor 已提交
183 184 185 186 187 188
    helper.append_op(
        type='arg_max', inputs={'X': x}, outputs={'Out': [out]}, attrs=attrs)
    out.stop_gradient = True
    return out


189
def argmin(x, axis=None, keepdim=False, dtype="int64", name=None):
W
wawltor 已提交
190 191 192 193 194 195 196 197 198 199
    """
    This OP computes the indices of the min elements of the input tensor's
    element along the provided axis.

    Args:
        x(Tensor): 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 x.ndim. when axis < 0, it works the same way
            as axis + R. Default is None, the input `x` will be into the flatten tensor, and selecting the min value index.
200
        keepdim(bool, optional): Keep the axis that selecting min. The defalut value is False.
W
wawltor 已提交
201
        dtype(str): Data type of the output tensor which can
202
                    be int32, int64. The default value is 'int64', and it will
W
wawltor 已提交
203 204 205 206 207 208 209 210 211 212 213 214 215
                    return the int64 indices.
        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:
        Tensor, return the tensor of `int32` if set :attr:`dtype` is `int32`, otherwise return the tensor of `int64`

    Examples:
        .. code-block:: python

            import paddle

216 217 218
            x =  paddle.to_tensor([[5,8,9,5],
                                     [0,0,1,7],
                                     [6,9,2,4]])
W
wawltor 已提交
219
            out1 = paddle.argmin(x)
N
Noel 已提交
220
            print(out1) # 4
W
wawltor 已提交
221
            out2 = paddle.argmin(x, axis=1)
N
Noel 已提交
222
            print(out2) 
W
wawltor 已提交
223 224
            # [0 0 2]
            out3 = paddle.argmin(x, axis=-1)
N
Noel 已提交
225
            print(out3) 
W
wawltor 已提交
226 227
            # [0 0 2]
    """
228 229 230 231
    if axis is not None and not isinstance(axis, int):
        raise TypeError(
            "The type of 'axis'  must be int or None in argmin, but received %s."
            % (type(axis)))
232

233 234 235 236
    if dtype is None:
        raise ValueError(
            "the value of 'dtype' in argmin could not be None, but received None"
        )
237

238 239
    var_dtype = convert_np_dtype_to_dtype_(dtype)
    check_dtype(var_dtype, 'dtype', ['int32', 'int64'], 'argmin')
W
wawltor 已提交
240
    flatten = False
241
    if axis is None:
W
wawltor 已提交
242 243 244 245
        flatten = True
        axis = 0

    if in_dygraph_mode():
246 247
        out = core.ops.arg_min(x, 'axis', axis, 'dtype', var_dtype, 'keepdims',
                               keepdim, 'flatten', flatten)
W
wawltor 已提交
248 249 250 251 252 253 254
        return out

    helper = LayerHelper("argmin", **locals())
    check_variable_and_dtype(
        x, 'x', ['float32', 'float64', 'int16', 'int32', 'int64', 'uint8'],
        'paddle.argmin')
    out = helper.create_variable_for_type_inference(var_dtype)
255
    attrs = {}
W
wawltor 已提交
256
    attrs['keepdims'] = keepdim
257
    attrs['axis'] = axis
W
wawltor 已提交
258
    attrs['flatten'] = flatten
259
    attrs['dtype'] = var_dtype
260
    helper.append_op(
W
wawltor 已提交
261
        type='arg_min', inputs={'X': x}, outputs={'Out': [out]}, attrs=attrs)
262 263
    out.stop_gradient = True
    return out
264 265


266
def index_select(x, index, axis=0, name=None):
267
    """
S
swtkiwi 已提交
268

269 270 271 272
    Returns a new tensor which indexes the ``input`` tensor along dimension ``axis`` using 
    the entries in ``index`` which is a Tensor. The returned tensor has the same number 
    of dimensions as the original ``x`` tensor. The dim-th dimension has the same 
    size as the length of ``index``; other dimensions have the same size as in the ``x`` tensor. 
C
Chengmo 已提交
273

274
    Args:
275 276 277
        x (Tensor): The input Tensor to be operated. The data of ``x`` can be one of float32, float64, int32, int64.
        index (Tensor): The 1-D Tensor containing the indices to index. The data type of ``index`` must be int32 or int64.
        axis (int, optional): The dimension in which we index. Default: if None, the ``axis`` is 0.
278 279 280
        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`.
281 282

    Returns:
283
        Tensor: A Tensor with same data type as ``x``.
284
    
285 286
    Examples:
        .. code-block:: python
287
            
288 289
            import paddle

290 291 292 293
            x = paddle.to_tensor([[1.0, 2.0, 3.0, 4.0],
                                  [5.0, 6.0, 7.0, 8.0],
                                  [9.0, 10.0, 11.0, 12.0]])
            index = paddle.to_tensor([0, 1, 1], dtype='int32')
294 295 296 297 298 299 300 301
            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.]]
302
    """
303

304
    if in_dygraph_mode():
305
        return core.ops.index_select(x, index, 'dim', axis)
306

307 308 309
    helper = LayerHelper("index_select", **locals())
    check_variable_and_dtype(x, 'x', ['float32', 'float64', 'int32', 'int64'],
                             'paddle.tensor.search.index_select')
310
    check_variable_and_dtype(index, 'index', ['int32', 'int64'],
311
                             'paddle.tensor.search.index_select')
312

313
    out = helper.create_variable_for_type_inference(x.dtype)
314 315 316

    helper.append_op(
        type='index_select',
317
        inputs={'X': x,
318 319
                'Index': index},
        outputs={'Out': out},
320
        attrs={'dim': axis})
321 322 323
    return out


324
def nonzero(x, as_tuple=False):
325 326 327 328 329 330 331 332
    """
    Return a tensor containing the indices of all non-zero elements of the `input` 
    tensor. If as_tuple is True, return a tuple of 1-D tensors, one for each dimension 
    in `input`, each containing the indices (in that dimension) of all non-zero elements 
    of `input`. Given a n-Dimensional `input` tensor with shape [x_1, x_2, ..., x_n], If 
    as_tuple is False, we can get a output tensor with shape [z, n], where `z` is the 
    number of all non-zero elements in the `input` tensor. If as_tuple is True, we can get 
    a 1-D tensor tuple of length `n`, and the shape of each 1-D tensor is [z, 1].
C
Chengmo 已提交
333

334
    Args:
335
        x (Tensor): The input tensor variable.
336 337 338
        as_tuple (bool): Return type, Tensor or tuple of Tensor.

    Returns:
339
        Tensor. The data type is int64.
340 341

    Examples:
342

N
Noel 已提交
343
        .. code-block:: python
李灿 已提交
344

345
            import paddle
346 347

            x1 = paddle.to_tensor([[1.0, 0.0, 0.0],
N
Noel 已提交
348 349
                                   [0.0, 2.0, 0.0],
                                   [0.0, 0.0, 3.0]])
350 351
            x2 = paddle.to_tensor([0.0, 1.0, 0.0, 3.0])
            out_z1 = paddle.nonzero(x1)
N
Noel 已提交
352
            print(out_z1)
353 354 355 356 357
            #[[0 0]
            # [1 1]
            # [2 2]]
            out_z1_tuple = paddle.nonzero(x1, as_tuple=True)
            for out in out_z1_tuple:
N
Noel 已提交
358
                print(out)
359 360 361 362 363 364 365
            #[[0]
            # [1]
            # [2]]
            #[[0]
            # [1]
            # [2]]
            out_z2 = paddle.nonzero(x2)
N
Noel 已提交
366
            print(out_z2)
367 368 369 370
            #[[1]
            # [3]]
            out_z2_tuple = paddle.nonzero(x2, as_tuple=True)
            for out in out_z2_tuple:
N
Noel 已提交
371
                print(out)
372 373
            #[[1]
            # [3]]
N
Noel 已提交
374

375 376
    """
    list_out = []
377
    shape = x.shape
378 379 380
    rank = len(shape)

    if in_dygraph_mode():
381
        outs = core.ops.where_index(x)
382
    else:
383
        outs = layers.where(x)
384 385 386 387 388 389 390 391 392

    if not as_tuple:
        return outs
    elif rank == 1:
        return tuple([outs])
    else:
        for i in range(rank):
            list_out.append(
                layers.slice(
393
                    outs, axes=[1], starts=[i], ends=[i + 1]))
394 395 396
        return tuple(list_out)


397
def sort(x, axis=-1, descending=False, name=None):
398
    """
S
swtkiwi 已提交
399

W
wawltor 已提交
400
    This OP sorts the input along the given axis, and returns the sorted output tensor. The default sort algorithm is ascending, if you want the sort algorithm to be descending, you must set the :attr:`descending` as True.
C
Chengmo 已提交
401

402
    Args:
403
        x(Tensor): An input N-D Tensor with type float32, float64, int16,
404 405 406 407 408 409 410 411 412 413 414
            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.
        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:
W
wawltor 已提交
415
        Tensor: sorted tensor(with the same shape and data type as ``x``).
416
    Examples:
N
Noel 已提交
417

418
        .. code-block:: python
N
Noel 已提交
419

420
            import paddle
N
Noel 已提交
421

422 423 424 425 426 427 428
            x = paddle.to_tensor([[[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')
429 430 431
            out1 = paddle.sort(x=x, axis=-1)
            out2 = paddle.sort(x=x, axis=0)
            out3 = paddle.sort(x=x, axis=1)
N
Noel 已提交
432
            print(out1)
W
wawltor 已提交
433 434 435 436 437 438
            #[[[5. 5. 8. 9.]
            #  [0. 0. 1. 7.]
            #  [2. 4. 6. 9.]]
            # [[2. 2. 4. 5.]
            #  [4. 7. 7. 9.]
            #  [0. 1. 6. 7.]]]
N
Noel 已提交
439
            print(out2)
440
            #[[[5. 2. 4. 2.]
W
wawltor 已提交
441 442 443 444 445
            #  [0. 0. 1. 7.]
            #  [1. 7. 0. 4.]]
            # [[5. 8. 9. 5.]
            #  [4. 7. 7. 9.]
            #  [6. 9. 2. 6.]]]
N
Noel 已提交
446
            print(out3)
447
            #[[[0. 0. 1. 4.]
W
wawltor 已提交
448 449 450 451 452
            #  [5. 8. 2. 5.]
            #  [6. 9. 9. 7.]]
            # [[1. 2. 0. 2.]
            #  [4. 7. 4. 6.]
            #  [5. 7. 7. 9.]]]
453
    """
454
    if in_dygraph_mode():
W
wawltor 已提交
455 456
        out, _ = core.ops.argsort(x, 'axis', axis, 'descending', descending)
        return out
457
    helper = LayerHelper("sort", **locals())
458 459
    out = helper.create_variable_for_type_inference(
        dtype=x.dtype, stop_gradient=False)
460 461 462 463
    ids = helper.create_variable_for_type_inference(
        VarDesc.VarType.INT64, stop_gradient=True)
    helper.append_op(
        type='argsort',
464
        inputs={'X': x},
465 466 467 468
        outputs={'Out': out,
                 'Indices': ids},
        attrs={'axis': axis,
               'descending': descending})
W
wawltor 已提交
469
    return out
C
Chengmo 已提交
470 471


472
def where(condition, x, y, name=None):
473
    r"""
474 475 476
    Return a tensor of elements selected from either $x$ or $y$, depending on $condition$.

    .. math::
C
Chengmo 已提交
477

478 479 480 481 482
      out_i =
      \\begin{cases}
      x_i, \quad  \\text{if}  \\ condition_i \\  is \\ True \\\\
      y_i, \quad  \\text{if}  \\ condition_i \\  is \\ False \\\\
      \\end{cases}
C
Chengmo 已提交
483

484

485
    Args:
G
GaoWei8 已提交
486 487 488
        condition(Tensor): The condition to choose x or y.
        x(Tensor): x is a Tensor with data type float32, float64, int32, int64.
        y(Tensor): y is a Tensor with data type float32, float64, int32, int64.
489 490 491 492 493

        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`.

494
    Returns:
G
GaoWei8 已提交
495
        Tensor: A Tensor with the same data dype as x. 
496

497 498 499
    Examples:
        .. code-block:: python

G
GaoWei8 已提交
500
          import paddle
501

502 503 504
          x = paddle.to_tensor([0.9383, 0.1983, 3.2, 1.2])
          y = paddle.to_tensor([1.0, 1.0, 1.0, 1.0])
          out = paddle.where(x>1, x, y)
505

G
GaoWei8 已提交
506
          print(out)
507
          #out: [1.0, 1.0, 3.2, 1.2]
508 509
    """
    if not in_dygraph_mode():
510
        check_variable_and_dtype(condition, 'condition', ['bool'], 'where')
511
        check_variable_and_dtype(
512
            x, 'x', ['float32', 'float64', 'int32', 'int64'], 'where')
513
        check_variable_and_dtype(
514
            y, 'y', ['float32', 'float64', 'int32', 'int64'], 'where')
515

516 517 518
    x_shape = list(x.shape)
    y_shape = list(y.shape)
    if x_shape == y_shape:
519
        if in_dygraph_mode():
520
            return core.ops.where(condition, x, y)
521 522
        else:
            helper = LayerHelper("where", **locals())
G
GaoWei8 已提交
523
            out = helper.create_variable_for_type_inference(dtype=x.dtype)
524 525 526

            helper.append_op(
                type='where',
527 528 529
                inputs={'Condition': condition,
                        'X': x,
                        'Y': y},
530 531 532
                outputs={'Out': [out]})
            return out
    else:
533 534 535 536
        cond_int = layers.cast(condition, x.dtype)
        cond_not_int = layers.cast(layers.logical_not(condition), x.dtype)
        out1 = layers.elementwise_mul(x, cond_int)
        out2 = layers.elementwise_mul(y, cond_not_int)
537 538 539 540
        out = layers.elementwise_add(out1, out2)
        return out


C
Chengmo 已提交
541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564
def index_sample(x, index):
    """
    **IndexSample Layer**

    IndexSample OP returns the element of the specified location of X, 
    and the location is specified by Index. 

    .. code-block:: text


                Given:

                X = [[1, 2, 3, 4, 5],
                     [6, 7, 8, 9, 10]]

                Index = [[0, 1, 3],
                         [0, 2, 4]]

                Then:

                Out = [[1, 2, 4],
                       [6, 8, 10]]

    Args:
C
Chengmo 已提交
565
        x (Tensor): The source input tensor with 2-D shape. Supported data type is 
C
Chengmo 已提交
566
            int32, int64, float32, float64.
C
Chengmo 已提交
567
        index (Tensor): The index input tensor with 2-D shape, first dimension should be same with X. 
C
Chengmo 已提交
568 569 570
            Data type is int32 or int64.

    Returns:
C
Chengmo 已提交
571
        output (Tensor): The output is a tensor with the same shape as index.
C
Chengmo 已提交
572 573 574 575 576 577

    Examples:

        .. code-block:: python

            import paddle
578 579 580 581 582 583 584 585 586 587 588

            x = paddle.to_tensor([[1.0, 2.0, 3.0, 4.0],
                                  [5.0, 6.0, 7.0, 8.0],
                                  [9.0, 10.0, 11.0, 12.0]], dtype='float32')
            index = paddle.to_tensor([[0, 1, 2],
                                      [1, 2, 3],
                                      [0, 0, 0]], dtype='int32')
            target = paddle.to_tensor([[100, 200, 300, 400],
                                       [500, 600, 700, 800],
                                       [900, 1000, 1100, 1200]], dtype='int32')
            out_z1 = paddle.index_sample(x, index)
N
Noel 已提交
589
            print(out_z1)
590 591 592 593 594 595 596 597
            #[[1. 2. 3.]
            # [6. 7. 8.]
            # [9. 9. 9.]]

            # Use the index of the maximum value by topk op
            # get the value of the element of the corresponding index in other tensors
            top_value, top_index = paddle.topk(x, k=2)
            out_z2 = paddle.index_sample(target, top_index)
N
Noel 已提交
598
            print(top_value)
599 600 601 602
            #[[ 4.  3.]
            # [ 8.  7.]
            # [12. 11.]]

N
Noel 已提交
603
            print(top_index)
604 605 606 607
            #[[3 2]
            # [3 2]
            # [3 2]]

N
Noel 已提交
608
            print(out_z2)
609 610 611
            #[[ 400  300]
            # [ 800  700]
            # [1200 1100]]
C
Chengmo 已提交
612

C
Chengmo 已提交
613
    """
C
Chengmo 已提交
614 615 616
    if in_dygraph_mode():
        return core.ops.index_sample(x, index)

C
Chengmo 已提交
617 618 619 620 621 622 623 624 625 626 627 628 629
    helper = LayerHelper("index_sample", **locals())
    check_variable_and_dtype(x, 'x', ['float32', 'float64', 'int32', 'int64'],
                             'paddle.tensor.search.index_sample')
    check_variable_and_dtype(index, 'index', ['int32', 'int64'],
                             'paddle.tensor.search.index_sample')
    out = helper.create_variable_for_type_inference(dtype=x.dtype)

    helper.append_op(
        type='index_sample',
        inputs={'X': x,
                'Index': index},
        outputs={'Out': out})
    return out
630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650


def masked_select(x, mask, name=None):
    """
    This OP Returns a new 1-D tensor which indexes the input tensor according to the ``mask``
    which is a tensor with data type of bool.

    Args:
        x (Tensor): The input Tensor, the data type can be int32, int64, float32, float64. 
        mask (Tensor): The Tensor containing the binary mask to index with, it's data type is bool.
        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: A 1-D Tensor which is the same data type  as ``x``.
    
    Examples:

        .. code-block:: python

            import paddle
651 652 653 654 655 656 657

            x = paddle.to_tensor([[1.0, 2.0, 3.0, 4.0],
                                  [5.0, 6.0, 7.0, 8.0],
                                  [9.0, 10.0, 11.0, 12.0]])
            mask = paddle.to_tensor([[True, False, False, False],
                                     [True, True, False, False],
                                     [True, False, False, False]])
658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674
            out = paddle.masked_select(x, mask)
            #[1.0 5.0 6.0 9.0]
    """

    if in_dygraph_mode():
        return core.ops.masked_select(x, mask)

    helper = LayerHelper("masked_select", **locals())
    check_variable_and_dtype(x, 'x', ['float32', 'float64', 'int32', 'int64'],
                             'paddle.tensor.search.mask_select')
    check_variable_and_dtype(mask, 'mask', ['bool'],
                             'paddle.tensor.search.masked_select')
    out = helper.create_variable_for_type_inference(dtype=x.dtype)
    helper.append_op(
        type='masked_select', inputs={'X': x,
                                      'Mask': mask}, outputs={'Y': out})
    return out
W
wawltor 已提交
675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703


def topk(x, k, axis=None, largest=True, sorted=True, name=None):
    """
    This OP is used to find values and indices of the k largest or smallest at the optional axis.
    If the input is a 1-D Tensor, finds the k largest or smallest values and indices.
    If the input is a Tensor with higher rank, this operator computes the top k values and indices along the :attr:`axis`.

    Args:
        x(Tensor): Tensor, an input N-D Tensor with type float32, float64, int32, int64.
        k(int, Tensor): The number of top elements to look for along the axis.
        axis(int, optional): Axis to compute indices along. The effective range
            is [-R, R), where R is x.ndim. when axis < 0, it works the same way
            as axis + R. Default is -1.
        largest(bool, optional) : largest is a flag, if set to true,
            algorithm will sort by descending order, otherwise sort by
            ascending order. Default is True.
        sorted(bool, optional): controls whether to return the elements in sorted order, default value is True. In gpu device, it always return the sorted value. 
        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        tuple(Tensor), return the values and indices. The value data type is the same as the input `x`. The indices data type is int64.

    Examples:

        .. code-block:: python

           import paddle

704
           tensor_1 = paddle.to_tensor([1, 4, 5, 7])
W
wawltor 已提交
705
           value_1, indices_1 = paddle.topk(tensor_1, k=1)
N
Noel 已提交
706
           print(value_1)
W
wawltor 已提交
707
           # [7]
N
Noel 已提交
708
           print(indices_1)
W
wawltor 已提交
709
           # [3] 
710
           tensor_2 = paddle.to_tensor([[1, 4, 5, 7], [2, 6, 2, 5]])
W
wawltor 已提交
711
           value_2, indices_2 = paddle.topk(tensor_2, k=1)
N
Noel 已提交
712
           print(value_2)
W
wawltor 已提交
713 714
           # [[7]
           #  [6]]
N
Noel 已提交
715
           print(indices_2)
W
wawltor 已提交
716 717 718
           # [[3]
           #  [1]]
           value_3, indices_3 = paddle.topk(tensor_2, k=1, axis=-1)
N
Noel 已提交
719
           print(value_3)
W
wawltor 已提交
720 721
           # [[7]
           #  [6]]
N
Noel 已提交
722
           print(indices_3)
W
wawltor 已提交
723 724 725
           # [[3]
           #  [1]]
           value_4, indices_4 = paddle.topk(tensor_2, k=1, axis=0)
N
Noel 已提交
726
           print(value_4)
W
wawltor 已提交
727
           # [[2 6 5 7]]
N
Noel 已提交
728
           print(indices_4)
W
wawltor 已提交
729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766
           # [[1 1 0 0]]

    """
    if in_dygraph_mode():
        k = k.numpy().item(0) if isinstance(k, Variable) else k
        if axis is None:
            out, indices = core.ops.top_k_v2(x, 'k',
                                             int(k), 'largest', largest,
                                             'sorted', sorted)
        else:
            out, indices = core.ops.top_k_v2(x, 'k',
                                             int(k), 'axis', axis, 'largest',
                                             largest, 'sorted', sorted)
        return out, indices

    helper = LayerHelper("top_k_v2", **locals())
    inputs = {"X": [x]}
    attrs = {}
    if isinstance(k, Variable):
        inputs['K'] = [k]
    else:
        attrs = {'k': k}
    attrs['largest'] = largest
    attrs['sorted'] = sorted
    if axis is not None:
        attrs['axis'] = axis

    values = helper.create_variable_for_type_inference(dtype=x.dtype)
    indices = helper.create_variable_for_type_inference(dtype="int64")

    helper.append_op(
        type="top_k_v2",
        inputs=inputs,
        outputs={"Out": [values],
                 "Indices": [indices]},
        attrs=attrs)
    indices.stop_gradient = True
    return values, indices