Function argpartition_device

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

Returns the indices that partition the array. Returns an error if the arguments are invalid.

The ordering of the elements within a partition in given by the indices is undefined.

Params

• a: The array to sort.
• kth: element index at the kth position in the output will give the sorted position. All indices before thekth position will be of elements less than or equal to the element at the kth index and all indices after will be elemenents greater than or equal to the element at the kth position.
• 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 = argpartition(&a, kth, axis);