eredu_core/
checkpoint.rs

1//! Neutral checkpoint tensor catalog contracts.
2
3use serde::{Deserialize, Serialize};
4use std::collections::BTreeMap;
5
6/// Portable tensor element type.
7#[derive(Debug, Clone, Eq, Hash, PartialEq, Serialize, Deserialize)]
8#[serde(rename_all = "snake_case")]
9pub enum TensorDtype {
10    /// Boolean.
11    Bool,
12    /// IEEE f32.
13    F32,
14    /// IEEE f16.
15    F16,
16    /// Brain float 16.
17    Bf16,
18    /// Signed 8-bit integer.
19    I8,
20    /// Unsigned 8-bit integer.
21    U8,
22    /// Unsigned 16-bit integer.
23    U16,
24    /// Unsigned 32-bit integer.
25    U32,
26    /// Unsigned 64-bit integer.
27    U64,
28    /// Signed 16-bit integer.
29    I16,
30    /// Signed 32-bit integer.
31    I32,
32    /// Signed 64-bit integer.
33    I64,
34    /// IEEE f64.
35    F64,
36    /// Complex number represented by two IEEE f32 values.
37    Complex64,
38    /// Backend-independent encoded/quantized storage.
39    Encoded(String),
40}
41
42/// Location of bytes within one checkpoint artifact member.
43#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
44pub struct TensorStorage {
45    /// Logical file/member name.
46    pub member: String,
47    /// Byte offset from the member start.
48    pub offset: u64,
49    /// Stored byte length.
50    pub length: u64,
51}
52
53/// Tensor descriptor without a materialized runtime array.
54#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
55pub struct TensorDescriptor {
56    /// Canonical checkpoint name.
57    pub name: String,
58    /// Row-major logical shape.
59    pub shape: Vec<usize>,
60    /// Logical or encoded dtype.
61    pub dtype: TensorDtype,
62    /// Optional source location.
63    pub storage: Option<TensorStorage>,
64}
65
66/// Validated name-indexed tensor catalog.
67#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
68pub struct TensorCatalog {
69    tensors: BTreeMap<String, TensorDescriptor>,
70}
71
72impl TensorCatalog {
73    /// Validates unique names and non-zero dimensions. An empty shape is a
74    /// valid rank-zero scalar with one element.
75    pub fn new(tensors: impl IntoIterator<Item = TensorDescriptor>) -> Result<Self, CatalogError> {
76        let mut map = BTreeMap::new();
77        for tensor in tensors {
78            if tensor.name.trim().is_empty() {
79                return Err(CatalogError::EmptyName);
80            }
81            if tensor.shape.contains(&0) {
82                return Err(CatalogError::InvalidShape(tensor.name));
83            }
84            let name = tensor.name.clone();
85            if map.insert(name.clone(), tensor).is_some() {
86                return Err(CatalogError::Duplicate(name));
87            }
88        }
89        Ok(Self { tensors: map })
90    }
91    /// Looks up a descriptor by canonical name.
92    pub fn get(&self, name: &str) -> Option<&TensorDescriptor> {
93        self.tensors.get(name)
94    }
95    /// Iterates over descriptors in deterministic name order.
96    pub fn descriptors(&self) -> impl Iterator<Item = &TensorDescriptor> {
97        self.tensors.values()
98    }
99    /// Number of cataloged tensors.
100    pub fn len(&self) -> usize {
101        self.tensors.len()
102    }
103    /// Whether the catalog is empty.
104    pub fn is_empty(&self) -> bool {
105        self.tensors.is_empty()
106    }
107}
108
109/// Tensor catalog validation error.
110#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
111pub enum CatalogError {
112    /// Tensor name is empty.
113    #[error("checkpoint tensor name must not be empty")]
114    EmptyName,
115    /// Tensor name is duplicated.
116    #[error("duplicate checkpoint tensor {0}")]
117    Duplicate(String),
118    /// Shape contains a zero dimension.
119    #[error("checkpoint tensor {0} has an invalid shape")]
120    InvalidShape(String),
121}
eredu_core/checkpoint.rs

eredu_core/
checkpoint.rs
