cherenkov/
tensors.rs

1use anyhow::{Result, bail};
2use cherenkov_model_data::{Checkpoint, Dtype as StoredDtype, ObjectId, TensorEncoding};
3use std::{collections::HashMap, path::Path};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum Dtype {
7    U32,
8    F32,
9    F16,
10    BF16,
11    I64,
12}
13
14impl Dtype {
15    pub fn size(self) -> usize {
16        match self {
17            Dtype::I64 => 8,
18            Dtype::U32 | Dtype::F32 => 4,
19            Dtype::F16 | Dtype::BF16 => 2,
20        }
21    }
22}
23
24#[derive(Debug, Clone)]
25pub struct TensorInfo {
26    pub dtype: Dtype,
27    pub shape: Vec<usize>,
28    /// Object index within the source checkpoint.
29    pub shard: usize,
30    /// Absolute byte offset of the tensor data within the shard file.
31    pub offset: usize,
32    pub nbytes: usize,
33}
34
35pub struct ModelWeights {
36    pub checkpoint: Checkpoint,
37    pub tensors: HashMap<String, TensorInfo>,
38}
39
40impl ModelWeights {
41    /// Keep the container's mappings; packing reads bytes without realigning
42    /// or duplicating the checkpoint in memory.
43    pub fn load_raw(path: &Path) -> Result<Self> {
44        let checkpoint = Checkpoint::open(path)?;
45        let mut tensors = HashMap::new();
46
47        for tensor in &checkpoint.inventory.tensors {
48            let TensorEncoding::Dense { tensor: stored } = &tensor.encoding else {
49                bail!(
50                    "{}: packer requires a supported dense storage type",
51                    tensor.name
52                );
53            };
54            let dtype = match stored.dtype {
55                StoredDtype::U32 => Dtype::U32,
56                StoredDtype::F32 => Dtype::F32,
57                StoredDtype::F16 => Dtype::F16,
58                StoredDtype::Bf16 => Dtype::BF16,
59                StoredDtype::I64 => Dtype::I64,
60                other => bail!("{}: packer does not support {other:?}", tensor.name),
61            };
62
63            tensors.insert(
64                tensor.name.clone(),
65                TensorInfo {
66                    dtype,
67                    shape: stored
68                        .shape
69                        .iter()
70                        .map(|&n| usize::try_from(n))
71                        .collect::<Result<_, _>>()?,
72                    shard: stored.data.object.0,
73                    offset: usize::try_from(stored.data.offset)?,
74                    nbytes: usize::try_from(stored.data.length)?,
75                },
76            );
77        }
78
79        Ok(Self {
80            checkpoint,
81            tensors,
82        })
83    }
84
85    pub fn tensor_bytes(&self, info: &TensorInfo) -> &[u8] {
86        // TensorInfo comes from validated container spans in load_raw.
87        let map = self
88            .checkpoint
89            .objects
90            .mapping(ObjectId(info.shard))
91            .expect("validated object");
92
93        &map[info.offset..info.offset + info.nbytes]
94    }
95}
cherenkov/tensors.rs

cherenkov/
tensors.rs
