mlx.core.partition

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

返回数组的一个分区副本，使得较小的 kth 元素排在前面。

分区中元素的顺序是未定义的。

参数:

• a (array) – 输入数组。

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

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

返回:

分区后的数组。

返回类型:

array