mlx.core.argmax

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

沿指定軸的最大值索引。

參數:

• a (array) -- 輸入陣列。

• axis (int, optional) -- 可選的歸約軸。若未指定，預設對整個陣列歸約。

• keepdims (bool, optional) -- 將被歸約的軸保留為單一維度，預設為 False。

回傳:

包含最大值索引的 uint32 陣列。

回傳型別:

array