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 pub fn has_nan_or_inf(&self) -> bool {
56 self.data.iter().any(|v| !v.is_finite())
57 }
58}
59
60#[cfg(test)]
61mod tests {
62 use super::*;
63
64 #[test]
65 fn row_indexing_matches_shape() {
66 let t = Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]);
67 assert_eq!(t.row(0), &[1.0, 2.0]);
68 assert_eq!(t.row(1), &[3.0, 4.0]);
69 assert_eq!(t.row(2), &[5.0, 6.0]);
70 }
71
72 #[test]
73 #[should_panic]
74 fn mismatched_shape_panics() {
75 Tensor::new(vec![1.0, 2.0, 3.0], vec![2, 2]);
76 }
77}