tenferro_ops/
axis.rs

1use thiserror::Error;
2
3/// Error returned when normalizing user-facing axis arguments.
4#[derive(Debug, Clone, PartialEq, Eq, Error)]
5pub enum AxisError {
6    /// Axis is outside `[-rank, rank)`.
7    #[error("axis {axis} is out of bounds for rank {rank}")]
8    OutOfBounds { axis: isize, rank: usize },
9    /// Axis appears more than once after negative-axis normalization.
10    #[error("duplicate axis {axis}")]
11    Duplicate { axis: usize },
12}
13
14/// Normalize a possibly-negative axis against `rank`.
15///
16/// # Examples
17///
18/// ```
19/// use tenferro_ops::axis::normalize_axis;
20///
21/// assert_eq!(normalize_axis(-1, 3).unwrap(), 2);
22/// assert!(normalize_axis(3, 3).is_err());
23/// ```
24pub fn normalize_axis(axis: isize, rank: usize) -> Result<usize, AxisError> {
25    let normalized = if axis >= 0 {
26        axis as usize
27    } else {
28        rank.checked_sub(axis.unsigned_abs())
29            .ok_or(AxisError::OutOfBounds { axis, rank })?
30    };
31    if normalized >= rank {
32        return Err(AxisError::OutOfBounds { axis, rank });
33    }
34    Ok(normalized)
35}
36
37/// Normalize a list of possibly-negative axes and reject duplicates.
38///
39/// # Examples
40///
41/// ```
42/// use tenferro_ops::axis::normalize_axes;
43///
44/// assert_eq!(normalize_axes(&[0, -1], 3).unwrap(), vec![0, 2]);
45/// assert!(normalize_axes(&[1, -2], 3).is_err());
46/// ```
47pub fn normalize_axes(axes: &[isize], rank: usize) -> Result<Vec<usize>, AxisError> {
48    let mut out = Vec::with_capacity(axes.len());
49    let mut seen = vec![false; rank];
50    for &axis in axes {
51        let normalized = normalize_axis(axis, rank)?;
52        if seen[normalized] {
53            return Err(AxisError::Duplicate { axis: normalized });
54        }
55        seen[normalized] = true;
56        out.push(normalized);
57    }
58    Ok(out)
59}
tenferro_ops/axis.rs

tenferro_ops/
axis.rs
