1#[derive(Debug, Clone, PartialEq)]
7pub struct Tensor {
8 pub data: Vec<f32>,
9 pub shape: Vec<usize>,
10}
11
12impl Tensor {
13 pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
14 let expected: usize = shape.iter().product();
15 assert_eq!(
16 data.len(),
17 expected,
18 "tensor data length {} does not match shape {:?} (expected {})",
19 data.len(),
20 shape,
21 expected
22 );
23 Tensor { data, shape }
24 }
25
26 pub fn zeros(shape: Vec<usize>) -> Self {
27 let n: usize = shape.iter().product();
28 Tensor {
29 data: vec![0.0; n],
30 shape,
31 }
32 }
33
34 pub fn len(&self) -> usize {
35 self.data.len()
36 }
37
38 pub fn is_empty(&self) -> bool {
39 self.data.is_empty()
40 }
41
42 pub fn rows(&self) -> usize {
43 self.shape.first().copied().unwrap_or(0)
44 }
45
46 pub fn cols(&self) -> usize {
47 self.shape.get(1).copied().unwrap_or(1)
48 }
49
50 pub fn row(&self, i: usize) -> &[f32] {
51 let cols = self.cols();
52 &self.data[i * cols..(i + 1) * cols]
53 }
54}
55
56#[cfg(test)]
57mod tests {
58 use super::*;
59
60 #[test]
61 fn row_indexing_matches_shape() {
62 let t = Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]);
63 assert_eq!(t.row(0), &[1.0, 2.0]);
64 assert_eq!(t.row(1), &[3.0, 4.0]);
65 assert_eq!(t.row(2), &[5.0, 6.0]);
66 }
67
68 #[test]
69 #[should_panic]
70 fn mismatched_shape_panics() {
71 Tensor::new(vec![1.0, 2.0, 3.0], vec![2, 2]);
72 }
73}