mlx.core.argpartition

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

返回对数组进行分区（partition）的索引。

通过这些索引获得的分区内元素的顺序是未定义的。

参数：

• a (array) – 输入数组。

• kth (int) – 输出中位于 kth 位置的元素的索引将给出其在排序后的位置。所有位于 kth 位置之前的索引对应的元素都小于或等于 kth 索引处对应的元素，所有位于 kth 索引之后的索引对应的元素都大于或等于 kth 索引处对应的元素。

• axis (int 或 None，可选) – 可选的用于分区的轴。如果为 None，则在展平的数组上进行分区。如果未指定，默认为 -1。

返回：

包含对输入进行分区（partition）的索引的 uint32 数组。

返回类型：

array