mindspore.mint.argmax
- mindspore.mint.argmax(input, dim=None, keepdim=False)[源代码]
返回输入Tensor在指定轴上的最大值索引。
- 参数:
input (Tensor) - 输入Tensor。
dim (Union[int, None],可选) - 指定计算轴。如果是
None
,将会返回扁平化Tensor在指定轴上的最大值索引。默认值:None
。keepdim (bool,可选) - 输出Tensor是否保留指定轴。如果 dim 是
None
,忽略该选项。默认值:False
。
- 返回:
Tensor,输出为指定轴上输入Tensor最大值的索引。
- 异常:
TypeError - 如果 keepdim 的类型不是bool值。
ValueError - 如果 dim 的设定值超出了范围。
- 支持平台:
Ascend
样例:
>>> import numpy as np >>> from mindspore import Tensor >>> from mindspore import mint >>> x = Tensor(np.array([[1, 20, 5], [67, 8, 9], [130, 24, 15]]).astype(np.float32)) >>> output = mint.argmax(x, dim=-1) >>> print(output) [1 0 0]