1use oxfuse_core::Error;
4use oxfuse_core::shape::checked_num_elements;
5
6#[derive(Debug, Clone, PartialEq)]
7pub struct Tensor {
8 shape: Vec<usize>,
9 data: Vec<f32>,
10}
11
12impl Tensor {
13 pub fn new(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, Error> {
14 Self::named("tensor", shape, data)
15 }
16
17 pub fn named(name: &str, shape: Vec<usize>, data: Vec<f32>) -> Result<Self, Error> {
19 let expected = checked_num_elements(&shape).ok_or_else(|| Error::InvalidShape {
20 name: name.to_owned(),
21 shape: shape.clone(),
22 reason: "element count overflows",
23 })?;
24 if data.len() != expected {
25 return Err(Error::DataLength {
26 name: name.to_owned(),
27 shape,
28 expected,
29 found: data.len(),
30 });
31 }
32 Ok(Self { shape, data })
33 }
34
35 pub fn shape(&self) -> &[usize] {
36 &self.shape
37 }
38
39 pub fn data(&self) -> &[f32] {
40 &self.data
41 }
42}
43
44#[cfg(test)]
45mod tests {
46 use super::*;
47
48 #[test]
49 fn new_accepts_matching_length() {
50 let t = Tensor::new(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
51 assert_eq!(t.shape(), &[2, 2]);
52 assert_eq!(t.data(), &[1.0, 2.0, 3.0, 4.0]);
53 }
54
55 #[test]
56 fn scalar_has_one_value() {
57 assert!(Tensor::new(vec![], vec![5.0]).is_ok());
58 assert!(Tensor::new(vec![], vec![]).is_err());
59 }
60
61 #[test]
62 fn named_error_carries_the_name() {
63 let err = Tensor::named("w", vec![2], vec![1.0]).unwrap_err();
64 assert_eq!(
65 err,
66 Error::DataLength {
67 name: "w".into(),
68 shape: vec![2],
69 expected: 2,
70 found: 1
71 }
72 );
73 }
74
75 #[test]
76 fn overflowing_shape_is_an_error_not_a_panic() {
77 let err = Tensor::new(vec![usize::MAX, 2], vec![]).unwrap_err();
78 assert!(matches!(err, Error::InvalidShape { .. }));
79 }
80}