mlx.core.partition#

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

Returns a partitioned copy of the array such that the smaller kth elements are first.

The ordering of the elements in partitions is undefined.

Parameters:

  • a (array) – Input array.

  • kth (int) – Element at the kth index will be in its sorted position in the output. All elements before the kth index will be less or equal to the kth element and all elements after will be greater or equal to the kth element in the output.

  • axis (int or None, optional) – Optional axis to partition over. If None, this partitions over the flattened array. If unspecified, it defaults to -1.

Returns:

The partitioned array.

Return type:

array