Skip to main content

topk

Function topk 

Source
pub fn topk(a: &Array, k: usize, axis: Option<isize>) -> Result<Array>
Expand description

The k largest elements of a along axis (flattened without one), in ascending order (MLX’s topk; MLX leaves the order unspecified).

use nuvec_core::{Array, Result, topk};

let a = Array::from_slice(&[5i32, 1, 9, 3, 7], &[5])?;
assert_eq!(topk(&a, 2, Some(0))?.to_vec::<i32>()?, [7, 9]);