mindspore.ops.argmax

mindspore.ops.argmax(input, dim=None, keepdim=False)

返回tensor在指定维度上的最大值索引。

参数：

• input (Tensor) - 输入tensor。

• dim (Union[int, None]，可选) - 指定计算的维度。如果为 None ，则计算 input 中所有元素的最大值索引。默认值： None 。

• keepdim (bool，可选) - 是否保留输出tensor的维度。默认值： False 。

返回：

Tensor，包含最大值的索引。

支持平台：

Ascend GPU CPU

样例：

>>> 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]])