cherenkov_model_data/
tensor.rs

1use anyhow::{Result, ensure};
2use serde::{Deserialize, Serialize};
3
4/// Identifies an object within an opened checkpoint. It is not a filename or
5/// a content digest; a store can assign content identities while copying it.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(transparent)]
8pub struct ObjectId(pub usize);
9
10/// A byte range within one object; offsets are independent of tensor encoding.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct DataSpan {
13    /// Object containing this range, relative to its byte source.
14    pub object: ObjectId,
15    /// Byte offset from the beginning of the object.
16    pub offset: u64,
17    /// Length of the range in bytes.
18    pub length: u64,
19}
20
21/// A logical tensor and the physical storage used to represent its values.
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct Tensor {
24    /// Original or adapter-assigned display name.
25    pub name: String,
26    /// Semantic role assigned by an architecture adapter.
27    pub role: TensorRole,
28    /// Logical dimensions, absent when an unknown encoding hides them.
29    pub logical: Option<TensorType>,
30    /// Physical representation and references to its stored bytes.
31    pub encoding: TensorEncoding,
32}
33
34/// Logical values are distinct from the physical types used to encode them.
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct TensorType {
37    /// Logical dimensions in axis order, before packing or quantization.
38    pub shape: Vec<u64>,
39    /// Dense tensors declare an element type. Quantized weights do not choose
40    /// the execution/accumulation type; a consumer must make that choice.
41    pub dtype: Option<Dtype>,
42}
43
44/// Index of a tensor in its containing inventory or model description.
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
46#[serde(transparent)]
47pub struct TensorId(pub usize);
48
49/// An architecture-assigned tensor role, independent of its checkpoint name.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
51#[serde(rename_all = "snake_case")]
52pub enum TensorRole {
53    Projection,
54    Embedding,
55    Expert,
56    Router,
57    Convolution,
58    Norm,
59    NgramEmbedding,
60    Buffer,
61    Opaque,
62}
63
64/// Scalar types supported by physical tensor views.
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
66#[serde(rename_all = "snake_case")]
67pub enum Dtype {
68    Bool,
69    U8,
70    I8,
71    U16,
72    I16,
73    I32,
74    U64,
75    F64,
76    U32,
77    I64,
78    F32,
79    F16,
80    Bf16,
81}
82
83impl Dtype {
84    /// Number of bytes occupied by one stored scalar.
85    pub fn bytes(self) -> u64 {
86        match self {
87            Self::I64 | Self::U64 | Self::F64 => 8,
88            Self::U32 | Self::I32 | Self::F32 => 4,
89            Self::F16 | Self::Bf16 | Self::U16 | Self::I16 => 2,
90            Self::Bool | Self::U8 | Self::I8 => 1,
91        }
92    }
93}
94
95/// Byte strides also describe interleaved records without introducing a
96/// table-specific layout or treating packed 4-bit values as addressable bytes.
97#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct StoredTensor {
99    /// Type of the addressable storage elements, such as packed U32 words.
100    pub dtype: Dtype,
101    /// Dimensions of the stored elements.
102    pub shape: Vec<u64>,
103    /// Byte distance between adjacent elements along each axis.
104    pub byte_strides: Vec<u64>,
105    /// Byte range containing the view, including any gaps between elements.
106    pub data: DataSpan,
107}
108
109impl StoredTensor {
110    /// Construct and validate a row-major view with the last axis contiguous.
111    /// Fail on size overflow or when the view exceeds its data span.
112    pub fn contiguous(dtype: Dtype, shape: Vec<u64>, data: DataSpan) -> Result<Self> {
113        let mut stride = dtype.bytes();
114        let mut byte_strides = vec![0; shape.len()];
115
116        for (dim, out) in shape.iter().zip(&mut byte_strides).rev() {
117            *out = stride;
118            stride = stride
119                .checked_mul(*dim)
120                .ok_or_else(|| anyhow::anyhow!("tensor size overflow"))?;
121        }
122
123        let tensor = Self {
124            dtype,
125            shape,
126            byte_strides,
127            data,
128        };
129
130        tensor.validate()?;
131
132        Ok(tensor)
133    }
134
135    /// Check stride rank, arithmetic overflow, and extent within the data span.
136    /// The backing object's bounds are checked separately by the byte source.
137    pub fn validate(&self) -> Result<()> {
138        ensure!(
139            self.shape.len() == self.byte_strides.len(),
140            "tensor stride rank mismatch"
141        );
142
143        ensure!(
144            self.data.offset.checked_add(self.data.length).is_some(),
145            "data span overflow"
146        );
147
148        if self.shape.contains(&0) {
149            return Ok(());
150        }
151
152        let mut extent = self.dtype.bytes();
153
154        for (&dim, &stride) in self.shape.iter().zip(&self.byte_strides) {
155            let last = (dim - 1)
156                .checked_mul(stride)
157                .ok_or_else(|| anyhow::anyhow!("tensor stride overflow"))?;
158            extent = extent
159                .checked_add(last)
160                .ok_or_else(|| anyhow::anyhow!("tensor extent overflow"))?;
161        }
162
163        ensure!(
164            extent <= self.data.length,
165            "tensor view exceeds its data span"
166        );
167
168        Ok(())
169    }
170}
171
172/// Storage layout and quantization metadata, with explicit references to all data.
173#[derive(Debug, Clone, Serialize, Deserialize)]
174#[serde(tag = "kind", rename_all = "snake_case")]
175pub enum TensorEncoding {
176    /// Addressable scalar values with explicit shape and strides.
177    Dense {
178        /// Physical view of the scalar values.
179        tensor: StoredTensor,
180    },
181    /// Grouped integer codes reconstructed using scales and an affine offset.
182    Affine {
183        /// Bits per quantized value.
184        bits: u8,
185        /// Logical values sharing one scale and offset.
186        group_size: u64,
187        /// Axis along which quantization groups are formed.
188        group_axis: usize,
189        /// Order of codes within each storage word.
190        packing: BitPacking,
191        /// Packed integer codes.
192        codes: StoredTensor,
193        /// One scale per quantization group.
194        scales: StoredTensor,
195        /// Additive biases or integer zero points, one per group.
196        offset: AffineOffset,
197    },
198    /// GGML blocks whose layout is determined by the encoding code.
199    Ggml {
200        encoding: GgmlEncoding,
201        data: DataSpan,
202    },
203    /// Uninterpreted storage retained for a future adapter.
204    Opaque {
205        name: String,
206        metadata: serde_json::Value,
207        data: Vec<DataSpan>,
208    },
209}
210
211/// Placement of quantized codes within addressable storage words.
212#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
213#[serde(rename_all = "snake_case")]
214pub enum BitPacking {
215    /// Consecutive codes occupy a U32 word from least to most significant bits.
216    LowFirstU32,
217}
218
219/// Per-group offset used to reconstruct affine-quantized values.
220#[derive(Debug, Clone, Serialize, Deserialize)]
221#[serde(tag = "kind", rename_all = "snake_case")]
222pub enum AffineOffset {
223    /// Reconstruction uses `code * scale + bias`.
224    Bias { tensor: StoredTensor },
225    /// Reconstruction uses `(code - zero_point) * scale`.
226    ZeroPoint { tensor: StoredTensor },
227}
228
229/// Recognized GGML element and block encodings; unknown numeric codes are preserved.
230#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
231#[serde(tag = "kind", rename_all = "snake_case")]
232pub enum GgmlEncoding {
233    F32,
234    F16,
235    Q4_0,
236    Q4_1,
237    Q4K,
238    Q5K,
239    Q6K,
240    Q8_0,
241    Bf16,
242    Opaque { code: u32 },
243}
244
245impl TensorEncoding {
246    /// Explicit references remain enumerable even for an unknown encoding.
247    pub fn data(&self) -> Vec<&DataSpan> {
248        match self {
249            Self::Dense { tensor } => vec![&tensor.data],
250            Self::Affine {
251                codes,
252                scales,
253                offset,
254                ..
255            } => {
256                let (AffineOffset::Bias { tensor } | AffineOffset::ZeroPoint { tensor }) = offset;
257
258                vec![&codes.data, &scales.data, &tensor.data]
259            }
260            Self::Ggml { data, .. } => vec![data],
261            Self::Opaque { data, .. } => data.iter().collect(),
262        }
263    }
264}
265
266impl Tensor {
267    /// Build a descriptor, inferring a logical dtype only for dense storage.
268    /// Call [`Self::validate`] to check consistency with the physical encoding.
269    pub fn new(
270        name: String,
271        role: TensorRole,
272        shape: Option<Vec<u64>>,
273        encoding: TensorEncoding,
274    ) -> Self {
275        let dtype = match &encoding {
276            TensorEncoding::Dense { tensor } => Some(tensor.dtype),
277            _ => None,
278        };
279
280        Self {
281            name,
282            role,
283            logical: shape.map(|shape| TensorType { shape, dtype }),
284            encoding,
285        }
286    }
287
288    /// Return logical dimensions, if the encoding's shape is known.
289    pub fn shape(&self) -> Option<&[u64]> {
290        self.logical.as_ref().map(|t| t.shape.as_slice())
291    }
292
293    /// Check the encoding and its agreement with the logical shape and dtype.
294    pub fn validate(&self) -> Result<()> {
295        self.encoding.validate()?;
296
297        if let TensorEncoding::Dense { tensor } = &self.encoding {
298            ensure!(
299                self.shape() == Some(tensor.shape.as_slice()),
300                "dense logical shape mismatch"
301            );
302            ensure!(
303                self.logical.as_ref().and_then(|t| t.dtype) == Some(tensor.dtype),
304                "dense logical dtype mismatch"
305            );
306        }
307
308        if let TensorEncoding::Affine {
309            codes,
310            bits,
311            group_axis,
312            ..
313        } = &self.encoding
314        {
315            let mut shape = codes.shape.clone();
316            shape[*group_axis] = shape[*group_axis]
317                .checked_mul(32 / u64::from(*bits))
318                .ok_or_else(|| anyhow::anyhow!("affine shape overflow"))?;
319
320            ensure!(
321                self.shape() == Some(shape.as_slice()),
322                "affine logical shape mismatch"
323            );
324        }
325
326        Ok(())
327    }
328}
329
330impl TensorEncoding {
331    /// Check component layouts and supported affine grouping constraints.
332    /// For GGML and opaque data, only span arithmetic is checked.
333    pub fn validate(&self) -> Result<()> {
334        match self {
335            Self::Dense { tensor } => tensor.validate(),
336            Self::Affine {
337                bits,
338                group_size,
339                group_axis,
340                codes,
341                scales,
342                offset,
343                ..
344            } => {
345                let (AffineOffset::Bias { tensor: offsets }
346                | AffineOffset::ZeroPoint { tensor: offsets }) = offset;
347
348                codes.validate()?;
349                scales.validate()?;
350                offsets.validate()?;
351                ensure!(
352                    matches!(bits, 2 | 4 | 8) && codes.dtype == Dtype::U32,
353                    "unsupported affine word packing"
354                );
355                ensure!(
356                    codes.shape.len().checked_sub(1) == Some(*group_axis),
357                    "affine groups must span the final axis"
358                );
359                ensure!(
360                    scales.shape.len() == codes.shape.len() && offsets.shape == scales.shape,
361                    "affine component rank mismatch"
362                );
363                ensure!(
364                    scales.shape[..*group_axis] == codes.shape[..*group_axis],
365                    "affine component row mismatch"
366                );
367
368                let width = codes.shape[*group_axis]
369                    .checked_mul(32 / u64::from(*bits))
370                    .ok_or_else(|| anyhow::anyhow!("affine shape overflow"))?;
371
372                ensure!(
373                    *group_size > 0 && width > 0 && width.is_multiple_of(*group_size),
374                    "invalid affine group size"
375                );
376                ensure!(
377                    scales.shape[*group_axis] == width / *group_size,
378                    "affine scale count mismatch"
379                );
380
381                Ok(())
382            }
383            Self::Ggml { data, .. } => check_span(data),
384            Self::Opaque { data, .. } => data.iter().try_for_each(check_span),
385        }
386    }
387}
388
389fn check_span(span: &DataSpan) -> Result<()> {
390    ensure!(
391        span.offset.checked_add(span.length).is_some(),
392        "data span overflow"
393    );
394
395    Ok(())
396}
cherenkov_model_data/tensor.rs

cherenkov_model_data/
tensor.rs
