burn_core/store/
mod.rs

1//! Minimal, non-generic record system for saving and loading module parameters.
2//!
3//! A [`ModuleRecord`](crate::store::ModuleRecord) holds a module's parameters (path +
4//! [`ParamId`](crate::module::ParamId) + [`TensorData`](crate::tensor::TensorData)) and
5//! serializes them with the [burnpack](burn_pack) format. It is produced and applied through the
6//! [`Module`](crate::module::Module) trait itself ([`Module::into_record`](crate::module::Module::into_record) /
7//! [`Module::load_record`](crate::module::Module::load_record)).
8//!
9//! This module is intentionally tiny: traversal is a straightforward
10//! [`ModuleVisitor`](crate::module::ModuleVisitor) / [`ModuleMapper`](crate::module::ModuleMapper)
11//! keyed by parameter path, with no filtering, adapters, or lazy snapshots.
12//! The richer snapshot/import tooling (filtering, key remapping, PyTorch/SafeTensors adapters,
13//! cross-framework stores) lives in the `burn-store` crate.
14
15use alloc::format;
16use alloc::string::{String, ToString};
17use alloc::vec::Vec;
18
19use hashbrown::HashMap;
20
21use crate::module::{Module, ModuleMapper, ModuleVisitor, Param, ParamId};
22use crate::tensor::{Bool, DType, Device, Float, Int, Shape, Tensor, TensorData, kind::Basic};
23
24use burn_pack::{Reader, Writer};
25
26/// Controls how a parameter's dtype is resolved when loading a [`ModuleRecord`].
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
28pub enum DTypePolicy {
29    /// The module parameter adopts the record's dtype (data is loaded verbatim). Default.
30    #[default]
31    FromRecord,
32    /// The record's data is cast to the module parameter's current dtype on load.
33    ///
34    /// Note this materializes each target parameter to read its dtype.
35    CastToModule,
36}
37
38/// Error returned by [`ModuleRecord`] save/load and [`Module`] apply operations.
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub enum RecordError {
41    /// An I/O or format error occurred while reading or writing the record.
42    Io(String),
43    /// Validation failed while applying the record (shape mismatch, or missing tensors
44    /// when partial loading is not allowed).
45    Validation(String),
46}
47
48impl core::fmt::Display for RecordError {
49    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
50        match self {
51            RecordError::Io(msg) => write!(f, "Record I/O error: {msg}"),
52            RecordError::Validation(msg) => write!(f, "Record validation error: {msg}"),
53        }
54    }
55}
56
57#[cfg(feature = "std")]
58impl std::error::Error for RecordError {}
59
60impl From<burn_pack::Error> for RecordError {
61    fn from(err: burn_pack::Error) -> Self {
62        RecordError::Io(err.to_string())
63    }
64}
65
66/// A single recorded tensor: its module path, parameter id, and data.
67#[derive(Clone)]
68struct RecordTensor {
69    path: String,
70    id: ParamId,
71    data: TensorData,
72}
73
74/// A non-generic record holding a module's parameters.
75///
76/// Obtain one from a module with [`Module::into_record`], then either save it
77/// ([`save`](ModuleRecord::save) / [`into_bytes`](ModuleRecord::into_bytes)) or apply it back with
78/// [`Module::load_record`]. Load-time behavior is
79/// configured with the builder methods; they are ignored when saving.
80///
81/// The save-side dtype is intentionally not configurable: use `module.cast(dtype)` before
82/// taking the record. The record stores whatever dtype the module currently holds.
83#[derive(Clone)]
84pub struct ModuleRecord {
85    tensors: Vec<RecordTensor>,
86    dtype_policy: DTypePolicy,
87    allow_partial: bool,
88    validate: bool,
89}
90
91impl core::fmt::Debug for ModuleRecord {
92    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
93        f.debug_struct("ModuleRecord")
94            .field("num_tensors", &self.tensors.len())
95            .field("dtype_policy", &self.dtype_policy)
96            .field("allow_partial", &self.allow_partial)
97            .field("validate", &self.validate)
98            .finish()
99    }
100}
101
102impl ModuleRecord {
103    fn from_tensors(tensors: Vec<RecordTensor>) -> Self {
104        Self {
105            tensors,
106            dtype_policy: DTypePolicy::default(),
107            allow_partial: false,
108            validate: true,
109        }
110    }
111
112    /// The number of tensors in the record.
113    pub fn len(&self) -> usize {
114        self.tensors.len()
115    }
116
117    /// Whether the record holds no tensors.
118    pub fn is_empty(&self) -> bool {
119        self.tensors.is_empty()
120    }
121
122    /// Set the dtype policy used when loading into a module.
123    pub fn with_dtype_policy(mut self, policy: DTypePolicy) -> Self {
124        self.dtype_policy = policy;
125        self
126    }
127
128    /// Cast the record's data to the module parameter dtypes on load.
129    ///
130    /// Sugar for [`with_dtype_policy(DTypePolicy::CastToModule)`](ModuleRecord::with_dtype_policy).
131    pub fn cast_to_module_dtype(self) -> Self {
132        self.with_dtype_policy(DTypePolicy::CastToModule)
133    }
134
135    /// Allow loading even when some module parameters are absent from the record.
136    pub fn allow_partial(mut self, allow: bool) -> Self {
137        self.allow_partial = allow;
138        self
139    }
140
141    /// Enable or disable validation while loading.
142    pub fn validate(mut self, validate: bool) -> Self {
143        self.validate = validate;
144        self
145    }
146
147    /// Serialize the record to an in-memory burnpack byte buffer.
148    pub fn into_bytes(self) -> Result<crate::tensor::Bytes, RecordError> {
149        Ok(Writer::new(self.pack_tensors()).into_bytes()?)
150    }
151
152    /// Reconstruct a record from an in-memory burnpack byte buffer.
153    pub fn from_bytes(bytes: crate::tensor::Bytes) -> Result<Self, RecordError> {
154        Self::from_reader(Reader::from_bytes(bytes)?)
155    }
156
157    /// Save the record to a burnpack file on disk.
158    #[cfg(feature = "std")]
159    pub fn save<P: AsRef<std::path::Path>>(self, path: P) -> Result<(), RecordError> {
160        Writer::new(self.pack_tensors()).write_to_file(path)?;
161        Ok(())
162    }
163
164    /// Load a record from a burnpack file on disk.
165    #[cfg(feature = "std")]
166    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self, RecordError> {
167        Self::from_reader(Reader::from_file(path)?)
168    }
169
170    fn pack_tensors(self) -> Vec<burn_pack::Tensor> {
171        self.tensors
172            .into_iter()
173            .map(|t| {
174                burn_pack::Tensor::new(
175                    t.path,
176                    t.data.dtype,
177                    t.data.shape,
178                    Some(t.id.val()),
179                    t.data.bytes,
180                )
181            })
182            .collect()
183    }
184
185    fn from_reader(reader: Reader) -> Result<Self, RecordError> {
186        let tensors = reader
187            .into_tensors()?
188            .into_iter()
189            .map(|t| {
190                let id = t.param_id.map(ParamId::from).unwrap_or_else(ParamId::new);
191                let data = TensorData::from_bytes(t.bytes, t.shape, t.dtype);
192                RecordTensor {
193                    path: t.name,
194                    id,
195                    data,
196                }
197            })
198            .collect();
199        Ok(Self::from_tensors(tensors))
200    }
201
202    /// Collect a module's parameters into a [`ModuleRecord`].
203    ///
204    /// Backs [`Module::into_record`](crate::module::Module::into_record).
205    pub(crate) fn from_module<M: Module>(module: M) -> Self {
206        let mut collector = Collector::default();
207        module.visit(&mut collector);
208        ModuleRecord::from_tensors(collector.tensors)
209    }
210
211    /// Apply this record to a module, returning the loaded module.
212    ///
213    /// Honors the record's [`DTypePolicy`], `validate`, and `allow_partial` settings. Backs
214    /// [`Module::try_load_record`](crate::module::Module::try_load_record).
215    pub(crate) fn apply<M: Module>(self, module: M) -> Result<M, RecordError> {
216        let validate = self.validate;
217        let allow_partial = self.allow_partial;
218
219        let mut mapper = ModuleRecordMapper::new(self);
220        let module = module.map(&mut mapper);
221
222        if validate && !mapper.errors.is_empty() {
223            return Err(RecordError::Validation(format!(
224                "Apply errors: {:?}",
225                mapper.errors
226            )));
227        }
228        if !allow_partial && !mapper.missing.is_empty() {
229            return Err(RecordError::Validation(format!(
230                "Missing tensors: {:?}",
231                mapper.missing
232            )));
233        }
234
235        Ok(module)
236    }
237}
238
239/// Visitor that collects every parameter as a [`RecordTensor`], keyed by its module path.
240#[derive(Default)]
241struct Collector {
242    path: Vec<String>,
243    tensors: Vec<RecordTensor>,
244}
245
246impl Collector {
247    fn record(&mut self, id: ParamId, data: TensorData) {
248        self.tensors.push(RecordTensor {
249            path: self.path.join("."),
250            id,
251            data,
252        });
253    }
254}
255
256impl ModuleVisitor for Collector {
257    fn enter_module(&mut self, name: &str, _container_type: &str) {
258        self.path.push(name.to_string());
259    }
260
261    fn exit_module(&mut self, _name: &str, _container_type: &str) {
262        self.path.pop();
263    }
264
265    // Record the `on_save` form: it is what the load side validates against and
266    // un-maps with `on_load`. Recording `val()` breaks any param whose mapper
267    // changes the shape (a `Col`-layout `Linear` weight).
268    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
269        self.record(param.id, param.transform_for_save().val().into_data());
270    }
271
272    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<D, Int>>) {
273        self.record(param.id, param.transform_for_save().val().into_data());
274    }
275
276    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<D, Bool>>) {
277        self.record(param.id, param.transform_for_save().val().into_data());
278    }
279}
280
281/// Mapper that loads recorded tensors back onto matching parameters by module path,
282/// restoring the persisted [`ParamId`] so optimizer state (keyed by id) survives
283/// save/load cycles.
284struct ModuleRecordMapper {
285    path: Vec<String>,
286    /// Map from module path to (persisted ParamId, tensor data).
287    tensors: HashMap<String, (ParamId, TensorData)>,
288    dtype_policy: DTypePolicy,
289    missing: Vec<String>,
290    errors: Vec<String>,
291}
292
293impl ModuleRecordMapper {
294    fn new(record: ModuleRecord) -> Self {
295        let tensors = record
296            .tensors
297            .into_iter()
298            .map(|t| (t.path, (t.id, t.data)))
299            .collect();
300        Self {
301            path: Vec::new(),
302            tensors,
303            dtype_policy: record.dtype_policy,
304            missing: Vec::new(),
305            errors: Vec::new(),
306        }
307    }
308
309    /// Look up the recorded tensor for the current path and build the tensor to load,
310    /// or `None` (recording it as missing / errored) to leave the parameter unchanged.
311    ///
312    /// Returns the tensor and the persisted [`ParamId`] on a hit.
313    ///
314    /// `module_dtype` is only evaluated on a hit under [`DTypePolicy::CastToModule`], so a
315    /// missing parameter never materializes its current value just to read a dtype.
316    fn take<const D: usize, K: Basic>(
317        &mut self,
318        device: &Device,
319        target_shape: Shape,
320        module_dtype: impl FnOnce() -> DType,
321    ) -> Option<(Tensor<D, K>, ParamId)> {
322        let path = self.path.join(".");
323        let (id, data) = match self.tensors.remove_entry(&path) {
324            Some(entry) => entry.1,
325            None => {
326                self.missing.push(path);
327                return None;
328            }
329        };
330
331        // Resolve the dtype to load with (CastToModule casts to the module's dtype).
332        let dtype = match self.dtype_policy {
333            DTypePolicy::FromRecord => data.dtype,
334            DTypePolicy::CastToModule => module_dtype(),
335        };
336
337        if data.shape != target_shape {
338            self.errors.push(format!(
339                "{path}: shape mismatch, expected {:?} but record has {:?}",
340                target_shape, data.shape
341            ));
342            return None;
343        }
344
345        Some((Tensor::from_data(data, (device, dtype)), id))
346    }
347}
348
349/// Generate a `ModuleMapper::map_*` method for one tensor kind. The three kinds differ only
350/// in the tensor type, so the body — collect identity, look the tensor up, load on a hit — is
351/// shared here.
352macro_rules! map_kind {
353    ($method:ident, $kind:ty) => {
354        fn $method<const D: usize>(
355            &mut self,
356            param: Param<Tensor<D, $kind>>,
357        ) -> Param<Tensor<D, $kind>> {
358            let device = param.lazy_device();
359            let shape = param.lazy_shape();
360            match self.take(&device, shape, || param.val().dtype()) {
361                Some((tensor, record_id)) => param.transform_for_load(tensor, record_id),
362                None => param,
363            }
364        }
365    };
366}
367
368impl ModuleMapper for ModuleRecordMapper {
369    fn enter_module(&mut self, name: &str, _container_type: &str) {
370        self.path.push(name.to_string());
371    }
372
373    fn exit_module(&mut self, _name: &str, _container_type: &str) {
374        self.path.pop();
375    }
376
377    map_kind!(map_float, Float);
378    map_kind!(map_int, Int);
379    map_kind!(map_bool, Bool);
380}
381
382#[cfg(all(test, feature = "std"))]
383mod tests {
384    use super::*;
385    use crate as burn;
386    use crate::module::{Module, Param};
387    use crate::tensor::Tensor;
388    use burn_tensor::Device;
389
390    #[derive(Module, Debug)]
391    struct Tiny {
392        weight: Param<Tensor<2>>,
393        bias: Param<Tensor<1>>,
394    }
395
396    impl Tiny {
397        fn new(weight: [[f32; 2]; 2], bias: [f32; 2], device: &Device) -> Self {
398            Self {
399                weight: Param::from_data(weight, device),
400                bias: Param::from_data(bias, device),
401            }
402        }
403    }
404
405    #[derive(Module, Debug)]
406    struct TinyWide {
407        weight: Param<Tensor<2>>,
408        bias: Param<Tensor<1>>,
409        gamma: Param<Tensor<1>>,
410    }
411
412    impl TinyWide {
413        fn zeros(device: &Device) -> Self {
414            Self {
415                weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device),
416                bias: Param::from_data([0.0, 0.0], device),
417                gamma: Param::from_data([0.0, 0.0], device),
418            }
419        }
420    }
421
422    fn weights(model: &Tiny) -> (Vec<f32>, Vec<f32>) {
423        (
424            model.weight.val().to_data().to_vec().unwrap(),
425            model.bias.val().to_data().to_vec().unwrap(),
426        )
427    }
428
429    #[test]
430    fn round_trip_in_memory() {
431        let device = Default::default();
432        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
433
434        let bytes = model.into_record().into_bytes().unwrap();
435        let record = ModuleRecord::from_bytes(bytes).unwrap();
436        assert_eq!(record.len(), 2);
437
438        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
439        let (w, b) = weights(&loaded);
440        assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]);
441        assert_eq!(b, vec![5.0, 6.0]);
442    }
443
444    #[test]
445    fn round_trip_file() {
446        let device = Default::default();
447        let dir = tempfile::tempdir().unwrap();
448        let path = dir.path().join("tiny.bpk");
449
450        Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device)
451            .into_record()
452            .save(&path)
453            .unwrap();
454
455        let record = ModuleRecord::load(&path).unwrap();
456        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
457        let (w, b) = weights(&loaded);
458        assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]);
459        assert_eq!(b, vec![5.0, 6.0]);
460    }
461
462    #[test]
463    fn missing_tensor_requires_allow_partial() {
464        let device = Default::default();
465        // Tiny has weight+bias; TinyWide also expects `gamma`, which the record lacks.
466        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
467
468        let strict = TinyWide::zeros(&device).try_load_record(record.clone());
469        assert!(matches!(strict, Err(RecordError::Validation(_))));
470
471        let partial = TinyWide::zeros(&device).try_load_record(record.allow_partial(true));
472        assert!(partial.is_ok());
473        let loaded = partial.unwrap();
474        // weight/bias were loaded; gamma kept its (zero) initialization.
475        assert_eq!(
476            loaded.weight.val().to_data().to_vec::<f32>().unwrap(),
477            vec![1.0, 2.0, 3.0, 4.0]
478        );
479        assert_eq!(
480            loaded.gamma.val().to_data().to_vec::<f32>().unwrap(),
481            vec![0.0, 0.0]
482        );
483    }
484
485    /// Build a `Tiny` whose parameters are zero-valued and carry the given dtype.
486    fn tiny_with_dtype(device: &Device, dtype: DType) -> Tiny {
487        Tiny {
488            weight: Param::from_tensor(
489                Tensor::<2>::from_data([[0.0, 0.0], [0.0, 0.0]], device).cast(dtype),
490            ),
491            bias: Param::from_tensor(Tensor::<1>::from_data([0.0, 0.0], device).cast(dtype)),
492        }
493    }
494
495    #[test]
496    fn dtype_policy_from_record_keeps_record_dtype() {
497        let device = Default::default();
498        // Record holds f32 data.
499        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
500
501        // Target params are f64, but the default policy (FromRecord) loads the data verbatim (f32).
502        let loaded = tiny_with_dtype(&device, DType::F64).load_record(record);
503        assert_eq!(loaded.weight.val().dtype(), DType::F32);
504        assert_eq!(loaded.bias.val().dtype(), DType::F32);
505    }
506
507    #[test]
508    fn dtype_policy_cast_to_module_uses_module_dtype() {
509        let device = Default::default();
510        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
511
512        // CastToModule casts the record's f32 data to the module's f64 params on load.
513        let loaded =
514            tiny_with_dtype(&device, DType::F64).load_record(record.cast_to_module_dtype());
515        assert_eq!(loaded.weight.val().dtype(), DType::F64);
516        assert_eq!(loaded.bias.val().dtype(), DType::F64);
517        // Values survive the cast.
518        assert_eq!(
519            loaded.weight.val().to_data().to_vec::<f64>().unwrap(),
520            vec![1.0, 2.0, 3.0, 4.0]
521        );
522    }
523
524    /// Build a `Tiny` whose `bias` shape (len 3) does not match the record's (len 2).
525    fn tiny_wrong_bias_shape(device: &Device) -> Tiny {
526        Tiny {
527            weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device),
528            bias: Param::from_data([0.0, 0.0, 0.0], device),
529        }
530    }
531
532    #[test]
533    fn shape_mismatch_fails_validation() {
534        let device = Default::default();
535        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
536
537        // `bias` shape mismatch is reported as a validation error by default.
538        let result = tiny_wrong_bias_shape(&device).try_load_record(record);
539        assert!(matches!(result, Err(RecordError::Validation(_))));
540    }
541
542    #[test]
543    fn load_record_preserves_param_id() {
544        let device = Default::default();
545        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
546
547        // Capture original ParamIds before saving.
548        let weight_id = model.weight.id;
549        let bias_id = model.bias.id;
550
551        let bytes = model.into_record().into_bytes().unwrap();
552        let record = ModuleRecord::from_bytes(bytes).unwrap();
553
554        // Load into a fresh model (new init() = different ParamIds).
555        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
556
557        // The loaded model should have the ORIGINAL ParamIds from the record,
558        // not the fresh ones from init().
559        assert_eq!(
560            loaded.weight.id, weight_id,
561            "weight ParamId should be restored from record"
562        );
563        assert_eq!(
564            loaded.bias.id, bias_id,
565            "bias ParamId should be restored from record"
566        );
567    }
568
569    #[test]
570    fn validate_false_ignores_shape_mismatch() {
571        let device = Default::default();
572        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
573
574        // With validation disabled, the shape-mismatched `bias` is skipped (keeps its value)
575        // while the matching `weight` still loads.
576        let loaded = tiny_wrong_bias_shape(&device)
577            .try_load_record(record.validate(false))
578            .unwrap();
579        assert_eq!(
580            loaded.weight.val().to_data().to_vec::<f32>().unwrap(),
581            vec![1.0, 2.0, 3.0, 4.0]
582        );
583        assert_eq!(
584            loaded.bias.val().to_data().to_vec::<f32>().unwrap(),
585            vec![0.0, 0.0, 0.0]
586        );
587    }
588
589    /// Mirrors a `Col`-layout `Linear` weight: persisted as `[3, 2]`, live as
590    /// the transposed `[2, 3]` through the init/save/load mappers.
591    #[derive(Module, Debug)]
592    struct ColLike {
593        weight: Param<Tensor<2>>,
594    }
595
596    impl ColLike {
597        fn new(seed: f32, device: &Device) -> Self {
598            let init_device = device.clone();
599            let weight = Param::uninitialized(
600                crate::module::ParamId::new(),
601                move |device, _| Tensor::<2>::full([3, 2], seed, device),
602                init_device,
603                true,
604                [3, 2].into(),
605            )
606            .init_mapper(|t: Tensor<2>| t.transpose())
607            .save_mapper(|t: Tensor<2>| t.transpose())
608            .load_mapper(|t: Tensor<2>| t.transpose());
609            Self { weight }
610        }
611    }
612
613    /// A param whose mapper changes the shape must round-trip through the
614    /// record in its save form.
615    #[test]
616    fn round_trip_a_shape_mapped_param() {
617        let device = Default::default();
618
619        let saved = ColLike::new(1.0, &device);
620        assert_eq!(saved.weight.val().dims(), [2, 3]);
621        let record = saved.into_record();
622        assert_eq!(
623            record.tensors[0].data.shape,
624            Shape::from([3, 2]),
625            "the record must hold the save form, not the live form"
626        );
627
628        let record = ModuleRecord::from_bytes(record.into_bytes().unwrap()).unwrap();
629        let loaded = ColLike::new(0.0, &device).load_record(record);
630        assert_eq!(loaded.weight.val().dims(), [2, 3]);
631        assert_eq!(
632            loaded.weight.val().to_data().to_vec::<f32>().unwrap(),
633            vec![1.0; 6],
634            "the recorded values must land, mapped back to the live form"
635        );
636    }
637}
burn_core/store/mod.rs

burn_core/store/
mod.rs
