Function partition_device

pub fn partition_device(
    a: impl AsRef<Array>,
    kth: i32,
    axis: i32,
    stream: impl AsRef<Stream>,
) -> Result<Array>

Returns a partitioned copy of the array such that the smaller kth elements are first. Returns an error if the arguments are invalid.

The ordering of the elements in partitions is undefined.

Params

• array: input array
• kth: 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: axis to partition over

Example

use mlx_rs::{Array, ops::*};

let a = Array::from_slice(&[3, 2, 1], &[3]);
let kth = 1;
let axis = 0;
let result = partition(&a, kth, axis);