oxmera_core/
shape.rs

1//! Tensor shapes and the broadcasting rules.
2
3use crate::error::{Error, Result};
4
5/// The extents of a tensor, one entry per dimension, outermost first
6/// (row-major convention throughout the project).
7///
8/// A rank-0 shape (`[]`) is a scalar and is valid. A dimension of size 0 is
9/// valid and makes the element count 0.
10#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub struct Shape(Vec<usize>);
12
13impl Shape {
14    /// Wrap dimension extents as a shape.
15    pub fn new(dims: Vec<usize>) -> Self {
16        Self(dims)
17    }
18
19    /// The dimension extents, outermost first.
20    pub fn dims(&self) -> &[usize] {
21        &self.0
22    }
23
24    /// The rank (number of dimensions).
25    pub fn ndim(&self) -> usize {
26        self.0.len()
27    }
28
29    /// The total number of elements.
30    ///
31    /// A scalar has 1 element; any zero-sized dimension makes this 0.
32    pub fn numel(&self) -> usize {
33        self.checked_numel().expect(
34            "shape element count overflows usize — reject the shape before it becomes a tensor",
35        )
36    }
37
38    /// Total number of elements, or `None` when the product of the
39    /// dimensions does not fit in `usize`.
40    ///
41    /// Every constructor that accepts a caller-supplied shape validates
42    /// with this before allocating: a wrapped product would let a shape
43    /// claim 2^64 elements over an empty buffer and pass the
44    /// `data.len() == numel` check (a reported class of bug); with the
45    /// check, such a shape is a typed error instead.
46    pub fn checked_numel(&self) -> Option<usize> {
47        self.0.iter().try_fold(1usize, |acc, &d| acc.checked_mul(d))
48    }
49}
50
51impl From<&[usize]> for Shape {
52    fn from(dims: &[usize]) -> Self {
53        Self(dims.to_vec())
54    }
55}
56
57impl<const N: usize> From<[usize; N]> for Shape {
58    fn from(dims: [usize; N]) -> Self {
59        Self(dims.to_vec())
60    }
61}
62
63impl From<Vec<usize>> for Shape {
64    fn from(dims: Vec<usize>) -> Self {
65        Self(dims)
66    }
67}
68
69/// The shape two operands broadcast to, or a typed error when they are
70/// incompatible.
71///
72/// The rules are NumPy's: align trailing dimensions; each pair must be
73/// equal or one of them 1. This is a total function over pairs of shapes —
74/// every input has a defined answer, success or a specific error.
75pub fn broadcast_shapes(lhs: &Shape, rhs: &Shape) -> Result<Shape> {
76    let (a, b) = (lhs.dims(), rhs.dims());
77    let ndim = a.len().max(b.len());
78    let mut out = vec![0usize; ndim];
79    for i in 0..ndim {
80        let da = if i < a.len() { a[a.len() - 1 - i] } else { 1 };
81        let db = if i < b.len() { b[b.len() - 1 - i] } else { 1 };
82        out[ndim - 1 - i] = if da == db {
83            da
84        } else if da == 1 {
85            db
86        } else if db == 1 {
87            da
88        } else {
89            return Err(Error::BroadcastIncompatible {
90                lhs: lhs.clone(),
91                rhs: rhs.clone(),
92            });
93        };
94    }
95    Ok(Shape(out))
96}
97
98#[cfg(test)]
99mod tests {
100    use super::Shape;
101
102    #[test]
103    fn checked_numel_matches_numel_when_it_fits() {
104        assert_eq!(Shape::from([2, 3, 4]).checked_numel(), Some(24));
105        assert_eq!(Shape::from([0, 3]).checked_numel(), Some(0));
106        assert_eq!(Shape::new(vec![]).checked_numel(), Some(1));
107    }
108
109    #[test]
110    fn checked_numel_is_none_on_overflow() {
111        assert_eq!(
112            Shape::from([1usize << 32, 1usize << 32]).checked_numel(),
113            None
114        );
115        assert_eq!(Shape::from([usize::MAX, 2]).checked_numel(), None);
116    }
117
118    #[test]
119    #[should_panic(expected = "overflows usize")]
120    fn numel_refuses_to_wrap_silently() {
121        let _ = Shape::from([1usize << 32, 1usize << 32]).numel();
122    }
123}
oxmera_core/shape.rs

oxmera_core/
shape.rs
