| //! Minimal row-major f32 CPU tensor: what `WeightStore::load_f32` returns (load-time conversions only; the data path is CUDA). | |
| pub struct Tensor { | |
| pub shape: Vec<usize>, | |
| pub data: Vec<f32>, | |
| } | |
| impl Tensor { | |
| pub fn new(shape: Vec<usize>, data: Vec<f32>) -> Self { | |
| assert_eq!(shape.iter().product::<usize>(), data.len(), "shape {:?} vs len {}", shape, data.len()); | |
| Tensor { shape, data } | |
| } | |
| pub fn numel(&self) -> usize { | |
| self.data.len() | |
| } | |
| /// rows of a matrix `[rows, cols]` (all leading dims flattened) | |
| pub fn rows(&self) -> usize { | |
| self.data.len() / self.cols() | |
| } | |
| pub fn cols(&self) -> usize { | |
| *self.shape.last().unwrap_or(&1) | |
| } | |
| pub fn row(&self, i: usize) -> &[f32] { | |
| let c = self.cols(); | |
| &self.data[i * c..(i + 1) * c] | |
| } | |
| } | |