mindspore.ops.argmax

mindspore.ops.argmax(input, dim=None, keepdim=False)[source]

Return the indices of the maximum values along a specified dimension of the tensor.

Parameters

input (Tensor) – The input tensor.
dim (Union[int, None], optional) – Specify the dimension for computation. If None , compute maximum value indexes of all elements in the input . Default None .
keepdim (bool, optional) – Whether the output tensor has dim retained. Default False .

Returns

Tensor, contains the index of the maximum value.

Supported Platforms:: Ascend GPU CPU

Examples

>>> import mindspore
>>> input = mindspore.tensor([[9, 3, 4, 5],
...                           [5, 2, 7, 4],
...                           [8, 1, 3, 6]])
>>> # case 1: By default, compute the maximum indice of all elements.
>>> mindspore.ops.argmax(input)
Tensor(shape=[], dtype=Int64, value= 0)
>>>
>>> # case 2: Compute maximum indice along dim 1.
>>> mindspore.ops.argmax(input, dim=1)
Tensor(shape=[3], dtype=Int64, value= [0, 2, 0])
>>>
>>> # case 3: If keepdim=True, the output shape will be same of that of the input.
>>> mindspore.ops.argmax(input, dim=1, keepdim=True)
Tensor(shape=[3, 1], dtype=Int64, value=
[[0],
 [2],
 [0]])