tenferro-internal-ops 0.1.0

Graph op vocabulary and AD rule implementations for the tenferro workspace.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

use thiserror::Error;

/// Error returned when normalizing user-facing axis arguments.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum AxisError {
    /// Axis is outside `[-rank, rank)`.
    #[error("axis {axis} is out of bounds for rank {rank}")]
    OutOfBounds { axis: isize, rank: usize },
    /// Axis appears more than once after negative-axis normalization.
    #[error("duplicate axis {axis}")]
    Duplicate { axis: usize },
}

/// Normalize a possibly-negative axis against `rank`.
///
/// # Examples
///
/// ```
/// use tenferro_ops::axis::normalize_axis;
///
/// assert_eq!(normalize_axis(-1, 3).unwrap(), 2);
/// assert!(normalize_axis(3, 3).is_err());
/// ```
pub fn normalize_axis(axis: isize, rank: usize) -> Result<usize, AxisError> {
    let rank_i = rank as isize;
    let normalized = if axis < 0 { rank_i + axis } else { axis };
    if normalized < 0 || normalized >= rank_i {
        return Err(AxisError::OutOfBounds { axis, rank });
    }
    Ok(normalized as usize)
}

/// Normalize a list of possibly-negative axes and reject duplicates.
///
/// # Examples
///
/// ```
/// use tenferro_ops::axis::normalize_axes;
///
/// assert_eq!(normalize_axes(&[0, -1], 3).unwrap(), vec![0, 2]);
/// assert!(normalize_axes(&[1, -2], 3).is_err());
/// ```
pub fn normalize_axes(axes: &[isize], rank: usize) -> Result<Vec<usize>, AxisError> {
    let mut out = Vec::with_capacity(axes.len());
    let mut seen = vec![false; rank];
    for &axis in axes {
        let normalized = normalize_axis(axis, rank)?;
        if seen[normalized] {
            return Err(AxisError::Duplicate { axis: normalized });
        }
        seen[normalized] = true;
        out.push(normalized);
    }
    Ok(out)
}