kime_model/
tensors.rs

1//! A set of named tensors over one blob of bytes, whichever format they came from.
2
3use std::collections::HashMap;
4
5use kime_tensor::{Blob, DType};
6
7use crate::error::{Error, Result};
8
9/// Where one tensor lives in its blob. Offsets are absolute and have been checked against the blob.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct Entry {
12    /// The tensor's name.
13    pub name: String,
14    /// Its element type.
15    pub dtype: DType,
16    /// Its shape, outermost first.
17    pub shape: Vec<usize>,
18    /// First byte.
19    pub start: usize,
20    /// One past the last byte.
21    pub end: usize,
22}
23
24impl Entry {
25    /// The number of elements.
26    #[must_use]
27    pub fn numel(&self) -> usize {
28        self.shape.iter().product()
29    }
30}
31
32/// A borrowed tensor.
33#[derive(Debug, Clone, Copy)]
34pub struct View<'a> {
35    /// The tensor's name.
36    pub name: &'a str,
37    /// Its element type.
38    pub dtype: DType,
39    /// Its shape.
40    pub shape: &'a [usize],
41    /// Its bytes, little endian, row major.
42    pub bytes: &'a [u8],
43}
44
45impl View<'_> {
46    /// The elements as f32, converting from f16 or bf16.
47    #[must_use]
48    pub fn to_f32(&self) -> Vec<f32> {
49        (0..self.bytes.len() / self.dtype.size())
50            .map(|i| self.dtype.read_f32(self.bytes, i))
51            .collect()
52    }
53}
54
55/// Named tensors backed by one blob, in the order their file lists them.
56#[derive(Debug)]
57pub struct Tensors {
58    blob: Blob,
59    entries: Vec<Entry>,
60    by_name: HashMap<String, usize>,
61}
62
63impl Tensors {
64    /// Builds the set from entries a parser has already checked against `blob`.
65    pub(crate) fn new(blob: Blob, entries: Vec<Entry>) -> Result<Self> {
66        let mut by_name = HashMap::with_capacity(entries.len());
67        for (i, e) in entries.iter().enumerate() {
68            debug_assert!(e.start <= e.end && e.end <= blob.len());
69            if by_name.insert(e.name.clone(), i).is_some() {
70                return Err(Error::format(format!("tensor {:?} appears twice", e.name)));
71            }
72        }
73        Ok(Self { blob, entries, by_name })
74    }
75
76    /// All entries, in file order.
77    #[must_use]
78    pub fn entries(&self) -> &[Entry] {
79        &self.entries
80    }
81
82    /// The index of a tensor by name.
83    #[must_use]
84    pub fn index(&self, name: &str) -> Option<usize> {
85        self.by_name.get(name).copied()
86    }
87
88    /// The tensor at `i`, in file order. Panics if `i` is out of range.
89    #[must_use]
90    pub fn view(&self, i: usize) -> View<'_> {
91        let e = &self.entries[i];
92        View { name: &e.name, dtype: e.dtype, shape: &e.shape, bytes: &self.blob[e.start..e.end] }
93    }
94
95    /// A tensor by name.
96    #[must_use]
97    pub fn get(&self, name: &str) -> Option<View<'_>> {
98        self.index(name).map(|i| self.view(i))
99    }
100
101    /// The bytes behind every tensor.
102    #[must_use]
103    pub fn blob(&self) -> &Blob {
104        &self.blob
105    }
106
107    /// Bytes of tensor data in total.
108    #[must_use]
109    pub fn data_bytes(&self) -> usize {
110        self.entries.iter().map(|e| e.end - e.start).sum()
111    }
112}
113
114/// The byte length of a tensor, or None when the shape overflows.
115pub(crate) fn byte_len(dtype: DType, shape: &[usize]) -> Option<usize> {
116    shape.iter().try_fold(dtype.size(), |acc, &d| acc.checked_mul(d))
117}
118
119/// Checks that no two ranges overlap. `ranges` holds (start, end, what) and is sorted in place.
120pub(crate) fn check_disjoint(ranges: &mut [(usize, usize, &str)]) -> Result<()> {
121    ranges.sort_unstable();
122    for w in ranges.windows(2) {
123        if w[1].0 < w[0].1 {
124            return Err(Error::format(format!("{:?} overlaps {:?}", w[1].2, w[0].2)));
125        }
126    }
127    Ok(())
128}
kime_model/tensors.rs

kime_model/
tensors.rs
