Skip to main content

oxfuse_runtime/
tensor.rs

1//! Owned, contiguous, row-major `f32` tensors.
2
3use 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    /// Like `new`, but errors name the tensor, for diagnostics on user data.
18    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}