nuvec_core/
axes.rs

1//! Axes and shapes as arguments: the sealed [`Axes`] and [`IntoShape`] traits, and the
2//! option structs of the reductions and scans that have rare knobs.
3
4use std::ops::{RangeFrom, RangeFull};
5
6use smallvec::SmallVec;
7
8use crate::shape::Shape;
9use crate::{Error, Result};
10
11mod sealed {
12    pub trait Sealed {}
13    impl<T: Sealed + ?Sized> Sealed for &T {}
14}
15
16/// `v` as an `i64` (saturating: an absurd `usize` fails validation, never wraps).
17pub(crate) fn wide(v: usize) -> i64 {
18    i64::try_from(v).unwrap_or(i64::MAX)
19}
20
21/// Which axes an operation works on, before it is resolved against an array.
22#[doc(hidden)]
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum AxesSpec {
25    /// Every axis (`..`, MLX's `axis=None`).
26    All,
27    /// The listed axes, negative counting from the end; an empty list means none.
28    List(SmallVec<[i64; 4]>),
29    /// Axes `k..ndim` (`k..`), `k` negative counting from the end.
30    From(i64),
31}
32
33/// Axes argument of reductions and shape ops: one axis (`-1`, `x.ndim() - 1`), a list
34/// (`[0, 2]`, `&axes[..]`, `vec![…]`, or a reference to any of these), all axes (`..`,
35/// MLX's `axis=None`) or the trailing axes from one on (`1..`).
36///
37/// An empty list means *no* axes, as in MLX C++ (`sum(a, {})`). Negative axes count from the
38/// end. Out-of-range or repeated axes make the operation fail with
39/// [`crate::OpErrorKind::InvalidAxis`].
40///
41/// ```
42/// use nuvec_core::Array;
43///
44/// let x = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [1, 2, 3]);
45/// assert_eq!(x.sum(-1).shape(), [1, 2]);
46/// assert_eq!(x.sum([0, 2]).shape(), [2]);
47/// assert_eq!(x.sum(..).shape(), [0i32; 0]);
48/// assert_eq!(x.sum(1..).shape(), [1]);
49/// assert_eq!(x.sum([]).shape(), [1, 2, 3]); // no axis reduced
50/// assert_eq!(x.sum(x.ndim() - 1).shape(), [1, 2]);
51/// ```
52#[diagnostic::on_unimplemented(
53    message = "`{Self}` is not an axes argument",
54    label = "expected an axis (`-1`), a list (`[0, 2]`, `&v[..]`), all axes (`..`) or trailing axes (`1..`)"
55)]
56pub trait Axes: sealed::Sealed {
57    #[doc(hidden)]
58    fn axes_spec(&self) -> AxesSpec;
59}
60
61#[diagnostic::do_not_recommend]
62impl<T: Axes + ?Sized> Axes for &T {
63    fn axes_spec(&self) -> AxesSpec {
64        (**self).axes_spec()
65    }
66}
67
68impl sealed::Sealed for i32 {}
69impl Axes for i32 {
70    fn axes_spec(&self) -> AxesSpec {
71        AxesSpec::List(SmallVec::from_elem(i64::from(*self), 1))
72    }
73}
74
75impl sealed::Sealed for usize {}
76impl Axes for usize {
77    fn axes_spec(&self) -> AxesSpec {
78        AxesSpec::List(SmallVec::from_elem(wide(*self), 1))
79    }
80}
81
82impl sealed::Sealed for [i32] {}
83impl Axes for [i32] {
84    fn axes_spec(&self) -> AxesSpec {
85        AxesSpec::List(self.iter().map(|&a| i64::from(a)).collect())
86    }
87}
88
89impl sealed::Sealed for [usize] {}
90impl Axes for [usize] {
91    fn axes_spec(&self) -> AxesSpec {
92        AxesSpec::List(self.iter().map(|&a| wide(a)).collect())
93    }
94}
95
96impl sealed::Sealed for Vec<i32> {}
97impl Axes for Vec<i32> {
98    fn axes_spec(&self) -> AxesSpec {
99        self.as_slice().axes_spec()
100    }
101}
102
103impl sealed::Sealed for Vec<usize> {}
104impl Axes for Vec<usize> {
105    fn axes_spec(&self) -> AxesSpec {
106        self.as_slice().axes_spec()
107    }
108}
109
110impl<const N: usize> sealed::Sealed for [i32; N] {}
111impl<const N: usize> Axes for [i32; N] {
112    fn axes_spec(&self) -> AxesSpec {
113        self.as_slice().axes_spec()
114    }
115}
116
117impl sealed::Sealed for RangeFull {}
118impl Axes for RangeFull {
119    fn axes_spec(&self) -> AxesSpec {
120        AxesSpec::All
121    }
122}
123
124impl sealed::Sealed for RangeFrom<i32> {}
125impl Axes for RangeFrom<i32> {
126    fn axes_spec(&self) -> AxesSpec {
127        AxesSpec::From(i64::from(self.start))
128    }
129}
130
131impl sealed::Sealed for RangeFrom<usize> {}
132impl Axes for RangeFrom<usize> {
133    fn axes_spec(&self) -> AxesSpec {
134        AxesSpec::From(wide(self.start))
135    }
136}
137
138/// The axis argument of operations that work along one axis (`argmax`, `sort`, the scans,
139/// `take`, `split`, `concatenate`, …): an `i32` or `usize` axis (negative counts from the
140/// end), or `..` for "the flattened array" where the operation has that meaning (as MLX's
141/// `axis=None`); elsewhere `..` fails with [`crate::OpErrorKind::InvalidAxis`].
142#[diagnostic::on_unimplemented(
143    message = "`{Self}` is not an axis argument",
144    label = "expected one axis (`-1`, `x.ndim() - 1`) or `..` (the flattened array)"
145)]
146pub trait Axis: sealed::Sealed {
147    /// The axis, or `None` for `..`.
148    #[doc(hidden)]
149    fn axis_spec(&self) -> Option<i64>;
150}
151
152#[diagnostic::do_not_recommend]
153impl<T: Axis + ?Sized> Axis for &T {
154    fn axis_spec(&self) -> Option<i64> {
155        (**self).axis_spec()
156    }
157}
158
159impl Axis for i32 {
160    fn axis_spec(&self) -> Option<i64> {
161        Some(i64::from(*self))
162    }
163}
164
165impl Axis for usize {
166    fn axis_spec(&self) -> Option<i64> {
167        Some(wide(*self))
168    }
169}
170
171impl Axis for RangeFull {
172    fn axis_spec(&self) -> Option<i64> {
173        None
174    }
175}
176
177/// One axis normalised against `ndim` dimensions, or `None` for `..` (flattened).
178pub(crate) fn one_axis(spec: Option<i64>, ndim: usize) -> Result<Option<usize>> {
179    spec.map(|a| normalize(a, ndim)).transpose()
180}
181
182/// An axis where `..` has no meaning: the axis as given (negative counts from the end, the
183/// operation normalises it), or an invalid-axis error for `..`.
184pub(crate) fn given_axis(spec: Option<i64>, ndim: usize) -> Result<isize> {
185    match spec {
186        Some(a) => Ok(a.clamp(isize::MIN as i64, isize::MAX as i64) as isize),
187        None => Err(Error::invalid_axis("axis", wide(ndim), ndim)),
188    }
189}
190
191/// `axis` as an index into `ndim` dimensions, or an invalid-axis error.
192pub(crate) fn normalize(axis: i64, ndim: usize) -> Result<usize> {
193    let n = ndim as i64;
194    let a = if axis < 0 { axis + n } else { axis };
195    if (0..n).contains(&a) {
196        Ok(a as usize)
197    } else {
198        Err(Error::invalid_axis("axis", axis, ndim))
199    }
200}
201
202impl AxesSpec {
203    /// The axes in the given order, normalised against `ndim`: every axis for `All`, `k..` for
204    /// `From(k)`. Out-of-range and repeated axes are errors.
205    pub(crate) fn list(&self, ndim: usize) -> Result<SmallVec<[usize; 4]>> {
206        match self {
207            Self::All => Ok((0..ndim).collect()),
208            Self::From(k) => {
209                // `ndim..` is valid and empty, like a slice.
210                let n = ndim as i64;
211                let start = if *k < 0 { k + n } else { *k };
212                if !(0..=n).contains(&start) {
213                    return Err(Error::invalid_axis("axis", *k, ndim));
214                }
215                Ok((start as usize..ndim).collect())
216            }
217            Self::List(axes) => {
218                let mut out = SmallVec::<[usize; 4]>::with_capacity(axes.len());
219                for &axis in axes {
220                    let a = normalize(axis, ndim)?;
221                    if out.contains(&a) {
222                        return Err(Error::invalid_axis("axis", axis, ndim));
223                    }
224                    out.push(a);
225                }
226                Ok(out)
227            }
228        }
229    }
230
231    /// The axes sorted (as reductions want them).
232    pub(crate) fn sorted(&self, ndim: usize) -> Result<SmallVec<[usize; 4]>> {
233        let mut axes = self.list(ndim)?;
234        axes.sort_unstable();
235        Ok(axes)
236    }
237}
238
239/// A shape argument: `[2, 3]` (of `i32` or `usize`), a slice, a vector, a reference to any of
240/// these, or `()` for a 0-D array. Every dimension must be in `0..=i32::MAX`; `reshape` also
241/// accepts one `-1` (inferred). Anything else makes the operation fail with
242/// [`crate::OpErrorKind::InvalidShape`].
243#[diagnostic::on_unimplemented(
244    message = "`{Self}` is not a shape argument",
245    label = "expected a shape such as `[2, 3]`, `&dims[..]`, `vec![2, 3]` or `()`"
246)]
247pub trait IntoShape: sealed::Sealed {
248    #[doc(hidden)]
249    fn shape_spec(&self) -> SmallVec<[i64; 4]>;
250}
251
252#[diagnostic::do_not_recommend]
253impl<T: IntoShape + ?Sized> IntoShape for &T {
254    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
255        (**self).shape_spec()
256    }
257}
258
259impl IntoShape for [i32] {
260    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
261        self.iter().map(|&d| i64::from(d)).collect()
262    }
263}
264
265impl IntoShape for [usize] {
266    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
267        self.iter().map(|&d| wide(d)).collect()
268    }
269}
270
271impl<const N: usize> IntoShape for [i32; N] {
272    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
273        self.as_slice().shape_spec()
274    }
275}
276
277impl<const N: usize> sealed::Sealed for [usize; N] {}
278impl<const N: usize> IntoShape for [usize; N] {
279    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
280        self.as_slice().shape_spec()
281    }
282}
283
284impl IntoShape for Vec<i32> {
285    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
286        self.as_slice().shape_spec()
287    }
288}
289
290impl IntoShape for Vec<usize> {
291    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
292        self.as_slice().shape_spec()
293    }
294}
295
296impl sealed::Sealed for () {}
297impl IntoShape for () {
298    fn shape_spec(&self) -> SmallVec<[i64; 4]> {
299        SmallVec::new()
300    }
301}
302
303/// A requested shape as dimension sizes: every dimension in `0..=i32::MAX`.
304pub(crate) fn dims(spec: &[i64]) -> Result<Shape> {
305    spec.iter()
306        .map(|&d| {
307            if (0..=i64::from(i32::MAX)).contains(&d) {
308                Ok(d as usize)
309            } else {
310                Err(Error::invalid_shape("shape", spec.iter().copied()))
311            }
312        })
313        .collect()
314}
315
316/// A requested `reshape` target for `size` elements: dimensions in `0..=i32::MAX` and at
317/// most one `-1`, inferred.
318pub(crate) fn reshape_dims(spec: &[i64], size: usize) -> Result<Shape> {
319    let invalid = || Error::invalid_shape("reshape", spec.iter().copied());
320    if spec.iter().any(|&d| d < -1 || d > i64::from(i32::MAX)) {
321        return Err(invalid());
322    }
323    let requested: Vec<isize> = spec.iter().map(|&d| d as isize).collect();
324    crate::view::infer_shape(size, &requested)
325}
326
327/// Rare knobs of [`crate::Array::var_with`] and [`crate::Array::std_with`].
328#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
329pub struct VarOptions {
330    /// Delta degrees of freedom: the divisor is `n - ddof` (clamped at 0). 0 by default.
331    pub ddof: i32,
332    /// Keep the reduced axes as dimensions of size 1.
333    pub keepdims: bool,
334}
335
336/// Rare knobs of the scans ([`crate::Array::cumsum_with`] and the other `*_with` scans).
337#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
338pub struct ScanOptions {
339    /// Scan from the end of the axis towards its start.
340    pub reverse: bool,
341    /// Leave each element out of its own output (MLX's `inclusive=False`).
342    pub exclusive: bool,
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348
349    #[test]
350    fn resolution() {
351        assert_eq!(AxesSpec::All.list(3).unwrap().as_slice(), [0, 1, 2]);
352        assert_eq!((-1).axes_spec().list(3).unwrap().as_slice(), [2]);
353        assert_eq!((1..).axes_spec().list(3).unwrap().as_slice(), [1, 2]);
354        assert_eq!((-1..).axes_spec().list(3).unwrap().as_slice(), [2]);
355        assert!((3..).axes_spec().list(3).unwrap().is_empty());
356        assert!((4..).axes_spec().list(3).is_err());
357        assert!([0, -3].axes_spec().list(3).is_err(), "repeated");
358        assert!(3.axes_spec().list(3).is_err());
359        assert_eq!([2, 0].axes_spec().sorted(3).unwrap().as_slice(), [0, 2]);
360        assert!(
361            <[i32; 0]>::default()
362                .axes_spec()
363                .list(3)
364                .unwrap()
365                .is_empty()
366        );
367        assert_eq!(one_axis((..).axis_spec(), 3).unwrap(), None);
368        assert_eq!(one_axis((-2).axis_spec(), 3).unwrap(), Some(1));
369        assert!(one_axis(usize::MAX.axis_spec(), 3).is_err(), "no wrap");
370        assert!(usize::MAX.axes_spec().list(3).is_err(), "no wrap");
371        assert!(given_axis((..).axis_spec(), 3).is_err());
372        assert_eq!((&&[0, 2]).axes_spec().list(3).unwrap().as_slice(), [0, 2]);
373        assert_eq!((1usize..).axes_spec().list(3).unwrap().as_slice(), [1, 2]);
374    }
375
376    #[test]
377    fn shapes() {
378        assert_eq!(dims(&[2, 3]).unwrap().as_slice(), [2, 3]);
379        assert!(dims(&[-1]).is_err());
380        assert!(dims(&[i64::from(i32::MAX) + 1]).is_err());
381        assert_eq!(reshape_dims(&[-1, 2], 6).unwrap().as_slice(), [3, 2]);
382        assert!(reshape_dims(&[-2, 3], 6).is_err());
383        assert!(reshape_dims(&[-1, -1], 6).is_err());
384        assert_eq!(().shape_spec().len(), 0);
385        assert_eq!([2usize, 3].shape_spec().as_slice(), [2, 3]);
386        let r: &[usize; 2] = &[2, 3];
387        assert_eq!(r.shape_spec().as_slice(), [2, 3]);
388        assert!(dims(&[usize::MAX].shape_spec()).is_err(), "no wrap");
389    }
390}
nuvec_core/axes.rs

nuvec_core/
axes.rs
