tenferro_ops/
reduction.rs

1use thiserror::Error;
2
3/// Error returned when computing public reduction output shapes.
4#[derive(Debug, Clone, PartialEq, Eq, Error)]
5pub enum ReductionShapeError {
6    /// Axis is outside the input rank.
7    #[error("axis {axis} is out of bounds for rank {rank}")]
8    AxisOutOfBounds { axis: usize, rank: usize },
9    /// Axis appears more than once.
10    #[error("duplicate axis {axis}")]
11    DuplicateAxis { axis: usize },
12}
13
14/// Compute the output shape for a reduction.
15///
16/// # Examples
17///
18/// ```
19/// use tenferro_ops::reduction::reduced_shape;
20///
21/// assert_eq!(reduced_shape(&[2, 3, 4], &[1], false).unwrap(), vec![2, 4]);
22/// assert_eq!(reduced_shape(&[2, 3, 4], &[1], true).unwrap(), vec![2, 1, 4]);
23/// ```
24///
25/// # Errors
26///
27/// Returns [`ReductionShapeError::AxisOutOfBounds`] for an invalid axis or
28/// [`ReductionShapeError::DuplicateAxis`] when an axis is repeated.
29pub fn reduced_shape(
30    input_shape: &[usize],
31    axes: &[usize],
32    keepdims: bool,
33) -> Result<Vec<usize>, ReductionShapeError> {
34    let mut reduced = vec![false; input_shape.len()];
35    for &axis in axes {
36        if axis >= input_shape.len() {
37            return Err(ReductionShapeError::AxisOutOfBounds {
38                axis,
39                rank: input_shape.len(),
40            });
41        }
42        if reduced[axis] {
43            return Err(ReductionShapeError::DuplicateAxis { axis });
44        }
45        reduced[axis] = true;
46    }
47    let mut out = Vec::with_capacity(input_shape.len());
48    for (axis, &dim) in input_shape.iter().enumerate() {
49        if reduced[axis] {
50            if keepdims {
51                out.push(1);
52            }
53        } else {
54            out.push(dim);
55        }
56    }
57    Ok(out)
58}
tenferro_ops/reduction.rs

tenferro_ops/
reduction.rs
