forge/
shape.rs

1use std::fmt;
2
3/// Row-major, contiguous tensor shape.
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub struct Shape(pub Vec<usize>);
6
7impl Shape {
8    pub fn new(dims: &[usize]) -> Self {
9        Shape(dims.to_vec())
10    }
11
12    pub fn dims(&self) -> &[usize] {
13        &self.0
14    }
15
16    pub fn rank(&self) -> usize {
17        self.0.len()
18    }
19
20    /// Total number of elements.
21    pub fn numel(&self) -> usize {
22        self.0.iter().product()
23    }
24
25    pub fn dim(&self, i: usize) -> usize {
26        self.0[i]
27    }
28
29    /// Row-major strides in elements.
30    pub fn strides(&self) -> Vec<usize> {
31        let mut strides = vec![1; self.0.len()];
32        for i in (0..self.0.len().saturating_sub(1)).rev() {
33            strides[i] = strides[i + 1] * self.0[i + 1];
34        }
35        strides
36    }
37}
38
39impl fmt::Display for Shape {
40    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41        write!(f, "{:?}", self.0)
42    }
43}
44
45impl From<&[usize]> for Shape {
46    fn from(d: &[usize]) -> Self {
47        Shape(d.to_vec())
48    }
49}
50
51impl From<Vec<usize>> for Shape {
52    fn from(d: Vec<usize>) -> Self {
53        Shape(d)
54    }
55}
56
57impl<const N: usize> From<[usize; N]> for Shape {
58    fn from(d: [usize; N]) -> Self {
59        Shape(d.to_vec())
60    }
61}
forge/shape.rs

forge/
shape.rs
