diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index fe2e979f9845c43ceb09f91bb7f3bc98059ad724..2c6a7f7ead1059763516ea11b9e41e28defe18d9 100644 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -123,7 +123,7 @@ def argmax(x, axis=None, keepdim=False, dtype="int64", name=None): 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. - keepdim(bool, optional): Keep the axis that selecting max. The defalut value is False. + keepdim(bool, optional): Whether to keep the given axis in output. If it is True, the dimensions will be same as input x and with size one in the axis. Otherwise the output dimentions is one fewer than x since the axis is squeezed. Default is False. 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. @@ -144,12 +144,15 @@ def argmax(x, axis=None, keepdim=False, dtype="int64", name=None): [6,9,2,4]]) out1 = paddle.argmax(x) print(out1) # 2 - out2 = paddle.argmax(x, axis=1) + out2 = paddle.argmax(x, axis=0) print(out2) - # [2 3 1] + # [2, 2, 0, 1] out3 = paddle.argmax(x, axis=-1) print(out3) - # [2 3 1] + # [2, 3, 1] + out4 = paddle.argmax(x, axis=0, keepdim=True) + print(out4) + # [[2, 2, 0, 1]] """ if axis is not None and not isinstance(axis, int): raise TypeError( @@ -200,7 +203,7 @@ def argmin(x, axis=None, keepdim=False, dtype="int64", name=None): 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. - keepdim(bool, optional): Keep the axis that selecting min. The defalut value is False. + keepdim(bool, optional): Whether to keep the given axis in output. If it is True, the dimensions will be same as input x and with size one in the axis. Otherwise the output dimentions is one fewer than x since the axis is squeezed. Default is False. dtype(str): Data type of the output tensor which can be int32, int64. The default value is 'int64', and it will return the int64 indices. @@ -221,12 +224,15 @@ def argmin(x, axis=None, keepdim=False, dtype="int64", name=None): [6,9,2,4]]) out1 = paddle.argmin(x) print(out1) # 4 - out2 = paddle.argmin(x, axis=1) + out2 = paddle.argmin(x, axis=0) print(out2) - # [0 0 2] + # [1, 1, 1, 2] out3 = paddle.argmin(x, axis=-1) print(out3) - # [0 0 2] + # [0, 0, 2] + out4 = paddle.argmin(x, axis=0, keepdim=True) + print(out4) + # [[1, 1, 1, 2]] """ if axis is not None and not isinstance(axis, int): raise TypeError(