nuvec_core/
sorting.rs

1//! Sorting and cumulative scans: `sort`, `argsort`, `partition`, `argpartition`, `topk`,
2//! `cumsum`, `cumprod`, `cummax`, `cummin` and `logcumsumexp`, as in MLX.
3//!
4//! Sorting is ascending and stable, with NaN last. `partition` is a full sort, which is a
5//! valid partition (MLX's GPU does the same; its CPU leaves the order unspecified).
6
7use crate::ops::{reshape, slice};
8use crate::primitive::{Primitive, ScanOp};
9use crate::view::normalize_axis;
10use crate::{Array, DType, Error, Result, SliceArg, at_least_float};
11
12fn single(primitive: Primitive, a: Array) -> Result<Array> {
13    Ok(Array::from_op(primitive, vec![a])?.swap_remove(0))
14}
15
16/// `a` and the axis to work on: `axis`, or 0 of the flattened array.
17fn along(a: &Array, axis: Option<isize>) -> Result<(Array, usize)> {
18    match axis {
19        Some(axis) => Ok((a.clone(), normalize_axis(axis, a.ndim())?)),
20        None => Ok((reshape(a, &[-1])?, 0)),
21    }
22}
23
24/// `a` sorted along `axis` (flattened without one): ascending, stable, NaN last.
25///
26/// ```
27/// use nuvec_core::{Array, Result, argsort, sort};
28///
29/// # fn main() -> Result<()> {
30/// let a = Array::from_slice(&[3.0f32, f32::NAN, 1.0, 2.0], &[4])?;
31/// assert_eq!(sort(&a, Some(0))?.to_vec::<f32>()?[..3], [1.0, 2.0, 3.0]);
32/// assert_eq!(argsort(&a, Some(0))?.to_vec::<u32>()?, [2, 3, 0, 1]);
33/// # Ok(())
34/// # }
35/// ```
36pub fn sort(a: &Array, axis: Option<isize>) -> Result<Array> {
37    let (a, axis) = along(a, axis)?;
38    single(Primitive::Sort { axis }, a)
39}
40
41/// The indices (uint32) that sort `a` along `axis`, as [`sort`].
42pub fn argsort(a: &Array, axis: Option<isize>) -> Result<Array> {
43    let (a, axis) = along(a, axis)?;
44    single(Primitive::ArgSort { axis }, a)
45}
46
47/// `kth` (negative counts from the end) validated against the axis length `n`.
48fn kth_index(op: &str, kth: isize, n: usize) -> Result<isize> {
49    let k = if kth < 0 { kth + n as isize } else { kth };
50    if (0..n as isize).contains(&k) {
51        Ok(k)
52    } else {
53        Err(Error::InvalidArgument(format!(
54            "{op}: kth {kth} is out of range for an axis of length {n}"
55        )))
56    }
57}
58
59/// `a` partitioned along `axis` (flattened without one): the element at position `kth`
60/// is the one a sort would put there, smaller ones come before it and larger ones after
61/// (MLX's `partition`).
62pub fn partition(a: &Array, kth: isize, axis: Option<isize>) -> Result<Array> {
63    let (a, axis) = along(a, axis)?;
64    kth_index("partition", kth, a.shape()[axis])?;
65    single(Primitive::Sort { axis }, a)
66}
67
68/// The indices (uint32) that partition `a`, as [`partition`].
69pub fn argpartition(a: &Array, kth: isize, axis: Option<isize>) -> Result<Array> {
70    let (a, axis) = along(a, axis)?;
71    kth_index("argpartition", kth, a.shape()[axis])?;
72    single(Primitive::ArgSort { axis }, a)
73}
74
75/// The `k` largest elements of `a` along `axis` (flattened without one), in ascending
76/// order (MLX's `topk`; MLX leaves the order unspecified).
77///
78/// ```
79/// use nuvec_core::{Array, Result, topk};
80///
81/// # fn main() -> Result<()> {
82/// let a = Array::from_slice(&[5i32, 1, 9, 3, 7], &[5])?;
83/// assert_eq!(topk(&a, 2, Some(0))?.to_vec::<i32>()?, [7, 9]);
84/// # Ok(())
85/// # }
86/// ```
87pub fn topk(a: &Array, k: usize, axis: Option<isize>) -> Result<Array> {
88    let (a, axis) = along(a, axis)?;
89    let n = a.shape()[axis];
90    if k > n {
91        return Err(Error::InvalidArgument(format!(
92            "topk: k = {k} exceeds the axis length {n}"
93        )));
94    }
95    if k == n {
96        return Ok(a);
97    }
98    let sorted = single(Primitive::Sort { axis }, a.clone())?;
99    let mut args = vec![SliceArg::full(); a.ndim()];
100    args[axis] = SliceArg::new(Some((n - k) as isize), None, 1);
101    slice(&sorted, &args)
102}
103
104fn scan(
105    op: ScanOp,
106    a: &Array,
107    axis: Option<isize>,
108    reverse: bool,
109    inclusive: bool,
110) -> Result<Array> {
111    let (a, axis) = along(a, axis)?;
112    let a = match op {
113        // As MLX: booleans sum as int32; log-sum-exp needs floating point.
114        ScanOp::Sum if a.dtype() == DType::Bool => a.astype(DType::I32)?,
115        ScanOp::LogAddExp => a.astype(at_least_float(a.dtype()))?,
116        _ => a,
117    };
118    // 8- and 16-bit integers scan in 32 bits and wrap back (the same modular results):
119    // the reverse exclusive scan of 8-bit integers was wrong on a virtualized Apple GPU
120    // (GitHub's macOS runners), while 32-bit scans are right everywhere.
121    let narrow = a.dtype();
122    let wide = match narrow {
123        DType::U8 | DType::U16 => Some(DType::U32),
124        DType::I8 | DType::I16 => Some(DType::I32),
125        _ => None,
126    };
127    let a = match wide {
128        Some(w) => a.astype(w)?,
129        None => a,
130    };
131    let out = single(
132        Primitive::Scan {
133            op,
134            axis,
135            reverse,
136            inclusive,
137        },
138        a,
139    )?;
140    match wide {
141        Some(_) => out.astype(narrow),
142        None => Ok(out),
143    }
144}
145
146/// The cumulative sum along `axis` (flattened without one): `reverse` scans from the
147/// end; an exclusive scan (`inclusive = false`) leaves each element out of its own
148/// output. Booleans sum as int32.
149///
150/// ```
151/// use nuvec_core::{Array, Result, cumsum};
152///
153/// # fn main() -> Result<()> {
154/// let a = Array::from_slice(&[1i32, 2, 3, 4], &[4])?;
155/// assert_eq!(cumsum(&a, None, false, true)?.to_vec::<i32>()?, [1, 3, 6, 10]);
156/// assert_eq!(cumsum(&a, None, true, false)?.to_vec::<i32>()?, [9, 7, 4, 0]);
157/// # Ok(())
158/// # }
159/// ```
160pub fn cumsum(a: &Array, axis: Option<isize>, reverse: bool, inclusive: bool) -> Result<Array> {
161    scan(ScanOp::Sum, a, axis, reverse, inclusive)
162}
163
164/// The cumulative product, as [`cumsum`].
165pub fn cumprod(a: &Array, axis: Option<isize>, reverse: bool, inclusive: bool) -> Result<Array> {
166    scan(ScanOp::Prod, a, axis, reverse, inclusive)
167}
168
169/// The running maximum (NaN propagates), as [`cumsum`].
170pub fn cummax(a: &Array, axis: Option<isize>, reverse: bool, inclusive: bool) -> Result<Array> {
171    scan(ScanOp::Max, a, axis, reverse, inclusive)
172}
173
174/// The running minimum (NaN propagates), as [`cumsum`].
175pub fn cummin(a: &Array, axis: Option<isize>, reverse: bool, inclusive: bool) -> Result<Array> {
176    scan(ScanOp::Min, a, axis, reverse, inclusive)
177}
178
179/// The running `log(Σ exp)`, as [`cumsum`] (computed in floating point).
180pub fn logcumsumexp(
181    a: &Array,
182    axis: Option<isize>,
183    reverse: bool,
184    inclusive: bool,
185) -> Result<Array> {
186    scan(ScanOp::LogAddExp, a, axis, reverse, inclusive)
187}
nuvec_core/sorting.rs

nuvec_core/
sorting.rs
