megengine.functional.argmax¶
- argmax(inp, axis=None, keepdims=False)[源代码]¶
返回在给定轴
axis
上inp
张量每一行的最大值的下标。如果给定的轴是一个维度列表,则规约所有维度。- 参数
- 返回类型
- 返回
输出张量。
实际案例
import numpy as np from megengine import tensor import megengine.functional as F x = tensor(np.arange(1, 7, dtype=np.int32).reshape(2,3)) out = F.argmax(x) print(out.numpy())
输出:
5