mlx.core.argmax

argmax(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: None | Stream | Device = None) → array

沿指定轴的最大值的索引。

参数:

• a (array) – 输入数组。

• axis (int, 可选) – 可选的轴，用于进行归约。如果未指定，则默认为对整个数组进行归约。

• keepdims (bool, 可选) – 是否保留被归约轴作为单维度，默认为 False。

返回:

包含最大值索引的 uint32 数组。

返回类型:

array