File size: 869 Bytes
19886eb ab0f9c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | //! Minimal row-major f32 CPU tensor: what `WeightStore::load_f32` returns (load-time conversions only; the data path is CUDA).
#[derive(Debug, Clone, PartialEq)]
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]
}
}
|