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
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}