mlx_rs/
options.rs

1/// Axis selection for operations that accept all, one, or several axes.
2///
3/// Convert an `i32`, `Vec<i32>`, slice, or array with `.into()`. Operations with additional
4/// independent defaults use a concrete `FooOptions: Default` type; for example,
5/// [`crate::fft::FftnOptions`] combines FFT lengths with this selection.
6#[derive(Debug, Clone, Default, PartialEq, Eq)]
7pub enum Axes {
8    /// Select every axis.
9    #[default]
10    All,
11
12    /// Select one axis.
13    Axis(i32),
14
15    /// Select several axes in the given order.
16    Axes(Vec<i32>),
17}
18
19impl From<i32> for Axes {
20    fn from(axis: i32) -> Self {
21        Self::Axis(axis)
22    }
23}
24
25impl From<Vec<i32>> for Axes {
26    fn from(axes: Vec<i32>) -> Self {
27        Self::Axes(axes)
28    }
29}
30
31impl From<&[i32]> for Axes {
32    fn from(axes: &[i32]) -> Self {
33        Self::Axes(axes.to_vec())
34    }
35}
36
37impl From<&Vec<i32>> for Axes {
38    fn from(axes: &Vec<i32>) -> Self {
39        Self::Axes(axes.clone())
40    }
41}
42
43impl<const N: usize> From<[i32; N]> for Axes {
44    fn from(axes: [i32; N]) -> Self {
45        Self::Axes(axes.into())
46    }
47}
48
49impl<const N: usize> From<&[i32; N]> for Axes {
50    fn from(axes: &[i32; N]) -> Self {
51        Self::Axes(axes.to_vec())
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[test]
60    fn ergonomic_conversions_preserve_axis_selection() {
61        assert_eq!(Axes::from(2), Axes::Axis(2));
62        assert_eq!(Axes::from(vec![0, -1]), Axes::Axes(vec![0, -1]));
63        assert_eq!(Axes::from(&vec![5, 6]), Axes::Axes(vec![5, 6]));
64        assert_eq!(Axes::from(&[1, 2]), Axes::Axes(vec![1, 2]));
65        assert_eq!(Axes::from(&[3, 4][..]), Axes::Axes(vec![3, 4]));
66    }
67}
mlx_rs/options.rs

mlx_rs/
options.rs
