Function argpartition_all

pub fn argpartition_all(a: impl AsRef<Array>, kth: i32) -> Result<Array>

Returns the indices that partition the flattened 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.

Example

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

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