Skip to main content

ferrox_core/
tensor.rs

1//! A minimal row-major dense f32 tensor. Ferrox keeps this deliberately
2//! small: weights live quantized in the mmap'd GGUF file and are
3//! dequantized on demand by ferrox-quant; this type is for activations
4//! and small dequantized weight slices during the forward pass.
5
6#[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}