mlx.core.topk

topk(a: array, /, k: int, axis: None | int = -1, *, stream: None | Stream | Device = None) → array

返回输入沿给定轴的 k 个最大元素。

这些元素不一定按排序顺序排列。

参数:

• a (array) – 输入数组。

• k (int) – 要返回的 top k 个元素

• axis (int 或 None，可选) – 可选的轴，用于选择。如果为 None，则选择展平数组中的 top k 个元素。如果未指定，则默认为 -1。

返回值:

输入中的 top k 个元素。

返回类型:

array