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, ParamGroup, 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` with `std` enabled / [`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    allow_unused: bool,
89    validate: bool,
90}
91
92impl core::fmt::Debug for ModuleRecord {
93    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
94        f.debug_struct("ModuleRecord")
95            .field("num_tensors", &self.tensors.len())
96            .field("dtype_policy", &self.dtype_policy)
97            .field("allow_partial", &self.allow_partial)
98            .field("allow_unused", &self.allow_unused)
99            .field("validate", &self.validate)
100            .finish()
101    }
102}
103
104impl ModuleRecord {
105    fn from_tensors(tensors: Vec<RecordTensor>) -> Self {
106        Self {
107            tensors,
108            dtype_policy: DTypePolicy::default(),
109            allow_partial: false,
110            allow_unused: false,
111            validate: true,
112        }
113    }
114
115    /// The number of tensors in the record.
116    pub fn len(&self) -> usize {
117        self.tensors.len()
118    }
119
120    /// Whether the record holds no tensors.
121    pub fn is_empty(&self) -> bool {
122        self.tensors.is_empty()
123    }
124
125    /// Set the dtype policy used when loading into a module.
126    pub fn with_dtype_policy(mut self, policy: DTypePolicy) -> Self {
127        self.dtype_policy = policy;
128        self
129    }
130
131    /// Cast the record's data to the module parameter dtypes on load.
132    ///
133    /// Sugar for [`with_dtype_policy(DTypePolicy::CastToModule)`](ModuleRecord::with_dtype_policy).
134    pub fn cast_to_module_dtype(self) -> Self {
135        self.with_dtype_policy(DTypePolicy::CastToModule)
136    }
137
138    /// Allow loading even when some module parameters are absent from the record.
139    ///
140    /// What a record of a part of a module needs — one taken with
141    /// [`into_record_group`](crate::module::Module::into_record_group), say — since it holds
142    /// nothing for the parameters outside that part.
143    pub fn allow_partial(mut self, allow: bool) -> Self {
144        self.allow_partial = allow;
145        self
146    }
147
148    /// Allow loading even when the record holds tensors no module parameter matches.
149    ///
150    /// The mirror of [`allow_partial`](Self::allow_partial), and refused by default: a record
151    /// entry that lands nowhere means the module is not the one the record was taken from, and
152    /// the load that looks like it succeeded has quietly done less than it was asked — or, at
153    /// `allow_partial(true)`, nothing at all. Allow it for the deliberate case, loading a
154    /// checkpoint into a part of the module it came from.
155    pub fn allow_unused(mut self, allow: bool) -> Self {
156        self.allow_unused = allow;
157        self
158    }
159
160    /// Enable or disable validation while loading.
161    pub fn validate(mut self, validate: bool) -> Self {
162        self.validate = validate;
163        self
164    }
165
166    /// Serialize the record to an in-memory burnpack byte buffer.
167    pub fn into_bytes(self) -> Result<crate::tensor::Bytes, RecordError> {
168        Ok(Writer::new(self.pack_tensors()).into_bytes()?)
169    }
170
171    /// Reconstruct a record from an in-memory burnpack byte buffer.
172    pub fn from_bytes(bytes: crate::tensor::Bytes) -> Result<Self, RecordError> {
173        Self::from_reader(Reader::from_bytes(bytes)?)
174    }
175
176    /// Save the record to a burnpack file on disk.
177    #[cfg(feature = "std")]
178    pub fn save<P: AsRef<std::path::Path>>(self, path: P) -> Result<(), RecordError> {
179        Writer::new(self.pack_tensors()).write_to_file(path)?;
180        Ok(())
181    }
182
183    /// Load a record from a burnpack file on disk.
184    #[cfg(feature = "std")]
185    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self, RecordError> {
186        Self::from_reader(Reader::from_file(path)?)
187    }
188
189    fn pack_tensors(self) -> Vec<burn_pack::Tensor> {
190        self.tensors
191            .into_iter()
192            .map(|t| {
193                burn_pack::Tensor::new(
194                    t.path,
195                    t.data.dtype,
196                    t.data.shape,
197                    Some(t.id.val()),
198                    t.data.bytes,
199                )
200            })
201            .collect()
202    }
203
204    fn from_reader(reader: Reader) -> Result<Self, RecordError> {
205        let tensors = reader
206            .into_tensors()?
207            .into_iter()
208            .map(|t| {
209                // Infallible for reader-produced tensors, which are always resident.
210                let (path, dtype, shape, param_id, bytes) = t.into_parts()?;
211                Ok(RecordTensor {
212                    path,
213                    id: param_id.map(ParamId::from).unwrap_or_else(ParamId::new),
214                    data: TensorData::from_bytes(bytes, shape, dtype),
215                })
216            })
217            .collect::<Result<Vec<_>, RecordError>>()?;
218        Ok(Self::from_tensors(tensors))
219    }
220
221    /// Collect a module's parameters into a [`ModuleRecord`].
222    ///
223    /// Backs [`Module::into_record`](crate::module::Module::into_record).
224    pub(crate) fn from_module<M: Module>(module: M, group: Option<ParamGroup>) -> Self {
225        let mut collector = Collector {
226            group,
227            ..Default::default()
228        };
229        module.visit(&mut collector);
230        ModuleRecord::from_tensors(collector.tensors)
231    }
232
233    /// Apply this record to a module, returning the loaded module.
234    ///
235    /// Honors the record's [`DTypePolicy`], `validate`, and `allow_partial` settings. Backs
236    /// [`Module::try_load_record`](crate::module::Module::try_load_record).
237    pub(crate) fn apply<M: Module>(self, module: M) -> Result<M, RecordError> {
238        let validate = self.validate;
239        let allow_partial = self.allow_partial;
240        let allow_unused = self.allow_unused;
241
242        let mut mapper = ModuleRecordMapper::new(self);
243        let module = module.map(&mut mapper);
244
245        if validate && !mapper.errors.is_empty() {
246            return Err(RecordError::Validation(format!(
247                "Apply errors: {:?}",
248                mapper.errors
249            )));
250        }
251        if !allow_partial && !mapper.missing.is_empty() {
252            return Err(RecordError::Validation(format!(
253                "Missing tensors: {:?}",
254                mapper.missing
255            )));
256        }
257        if !allow_unused && !mapper.unused().is_empty() {
258            return Err(RecordError::Validation(format!(
259                "Unused tensors: {:?}",
260                mapper.unused()
261            )));
262        }
263
264        Ok(module)
265    }
266}
267
268/// Visitor that collects a module's parameters as [`RecordTensor`]s, keyed by module path.
269///
270/// A [`ParamGroup`] narrows what is collected: a parameter the group does not
271/// match is skipped before its data is read, so a record of one group never
272/// materializes the rest of the module.
273#[derive(Default)]
274struct Collector {
275    path: Vec<String>,
276    group: Option<ParamGroup>,
277    tensors: Vec<RecordTensor>,
278}
279
280impl Collector {
281    fn record(&mut self, id: ParamId, data: impl FnOnce() -> TensorData) {
282        let path = self.path.join(".");
283        if let Some(group) = &self.group
284            && !group.matches(&id, Some(&path))
285        {
286            return;
287        }
288
289        self.tensors.push(RecordTensor {
290            path,
291            id,
292            data: data(),
293        });
294    }
295}
296
297impl ModuleVisitor for Collector {
298    fn enter_module(&mut self, name: &str, _container_type: &str) {
299        self.path.push(name.to_string());
300    }
301
302    fn exit_module(&mut self, _name: &str, _container_type: &str) {
303        self.path.pop();
304    }
305
306    // Record the `on_save` form: it is what the load side validates against and
307    // un-maps with `on_load`. Recording `val()` breaks any param whose mapper
308    // changes the shape (a `Col`-layout `Linear` weight).
309    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
310        self.record(param.id, || param.transform_for_save().val().into_data());
311    }
312
313    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<D, Int>>) {
314        self.record(param.id, || param.transform_for_save().val().into_data());
315    }
316
317    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<D, Bool>>) {
318        self.record(param.id, || param.transform_for_save().val().into_data());
319    }
320}
321
322/// Mapper that loads recorded tensors back onto matching parameters by module path,
323/// restoring the persisted [`ParamId`] so optimizer state (keyed by id) survives
324/// save/load cycles.
325struct ModuleRecordMapper {
326    path: Vec<String>,
327    /// Map from module path to (persisted ParamId, tensor data).
328    tensors: HashMap<String, (ParamId, TensorData)>,
329    dtype_policy: DTypePolicy,
330    missing: Vec<String>,
331    errors: Vec<String>,
332}
333
334impl ModuleRecordMapper {
335    fn new(record: ModuleRecord) -> Self {
336        let tensors = record
337            .tensors
338            .into_iter()
339            .map(|t| (t.path, (t.id, t.data)))
340            .collect();
341        Self {
342            path: Vec::new(),
343            tensors,
344            dtype_policy: record.dtype_policy,
345            missing: Vec::new(),
346            errors: Vec::new(),
347        }
348    }
349
350    /// The recorded tensors no parameter matched — what is left once the traversal has taken
351    /// every hit out. Sorted, so a message naming them reads the same twice.
352    fn unused(&self) -> Vec<&str> {
353        let mut unused: Vec<&str> = self.tensors.keys().map(String::as_str).collect();
354        unused.sort_unstable();
355        unused
356    }
357
358    /// Look up the recorded tensor for the current path and build the tensor to load,
359    /// or `None` (recording it as missing / errored) to leave the parameter unchanged.
360    ///
361    /// Returns the tensor and the persisted [`ParamId`] on a hit.
362    ///
363    /// `module_dtype` is only evaluated on a hit under [`DTypePolicy::CastToModule`], so a
364    /// missing parameter never materializes its current value just to read a dtype.
365    fn take<const D: usize, K: Basic>(
366        &mut self,
367        device: &Device,
368        target_shape: Shape,
369        module_dtype: impl FnOnce() -> DType,
370    ) -> Option<(Tensor<D, K>, ParamId)> {
371        let path = self.path.join(".");
372        let (id, data) = match self.tensors.remove_entry(&path) {
373            Some(entry) => entry.1,
374            None => {
375                self.missing.push(path);
376                return None;
377            }
378        };
379
380        // Resolve the dtype to load with (CastToModule casts to the module's dtype).
381        let dtype = match self.dtype_policy {
382            DTypePolicy::FromRecord => data.dtype,
383            DTypePolicy::CastToModule => module_dtype(),
384        };
385
386        if data.shape != target_shape {
387            self.errors.push(format!(
388                "{path}: shape mismatch, expected {:?} but record has {:?}",
389                target_shape, data.shape
390            ));
391            return None;
392        }
393
394        Some((Tensor::from_data(data, (device, dtype)), id))
395    }
396}
397
398/// Generate a `ModuleMapper::map_*` method for one tensor kind. The three kinds differ only
399/// in the tensor type, so the body — collect identity, look the tensor up, load on a hit — is
400/// shared here.
401macro_rules! map_kind {
402    ($method:ident, $kind:ty) => {
403        fn $method<const D: usize>(
404            &mut self,
405            param: Param<Tensor<D, $kind>>,
406        ) -> Param<Tensor<D, $kind>> {
407            let device = param.lazy_device();
408            let shape = param.lazy_shape();
409            match self.take(&device, shape, || param.val().dtype()) {
410                Some((tensor, record_id)) => param.transform_for_load(tensor, record_id),
411                None => param,
412            }
413        }
414    };
415}
416
417impl ModuleMapper for ModuleRecordMapper {
418    fn enter_module(&mut self, name: &str, _container_type: &str) {
419        self.path.push(name.to_string());
420    }
421
422    fn exit_module(&mut self, _name: &str, _container_type: &str) {
423        self.path.pop();
424    }
425
426    map_kind!(map_float, Float);
427    map_kind!(map_int, Int);
428    map_kind!(map_bool, Bool);
429}
430
431#[cfg(all(test, feature = "std"))]
432mod tests {
433    use super::*;
434    use crate as burn;
435    use crate::module::{Flag, Module, Param, ParamGroup};
436    use crate::tensor::Tensor;
437    use burn_tensor::Device;
438
439    #[derive(Module, Debug)]
440    struct Tiny {
441        weight: Param<Tensor<2>>,
442        bias: Param<Tensor<1>>,
443    }
444
445    impl Tiny {
446        fn new(weight: [[f32; 2]; 2], bias: [f32; 2], device: &Device) -> Self {
447            Self {
448                weight: Param::from_data(weight, device),
449                bias: Param::from_data(bias, device),
450            }
451        }
452    }
453
454    #[derive(Module, Debug)]
455    struct Flagged {
456        weight: Param<Tensor<1>>,
457        enabled: Param<Flag>,
458    }
459
460    impl Flagged {
461        fn new(weight: f32, enabled: bool, device: &Device) -> Self {
462            Self {
463                weight: Param::from_data([weight], device),
464                enabled: Param::from_bool(enabled),
465            }
466        }
467    }
468
469    #[derive(Module, Debug)]
470    struct TinyWide {
471        weight: Param<Tensor<2>>,
472        bias: Param<Tensor<1>>,
473        gamma: Param<Tensor<1>>,
474    }
475
476    impl TinyWide {
477        fn zeros(device: &Device) -> Self {
478            Self {
479                weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device),
480                bias: Param::from_data([0.0, 0.0], device),
481                gamma: Param::from_data([0.0, 0.0], device),
482            }
483        }
484    }
485
486    fn weights(model: &Tiny) -> (Vec<f32>, Vec<f32>) {
487        (
488            model.weight.val().try_into_vec_as().unwrap(),
489            model.bias.val().try_into_vec_as().unwrap(),
490        )
491    }
492
493    #[test]
494    fn round_trip_in_memory() {
495        let device = Default::default();
496        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
497
498        let bytes = model.into_record().into_bytes().unwrap();
499        let record = ModuleRecord::from_bytes(bytes).unwrap();
500        assert_eq!(record.len(), 2);
501
502        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
503        let (w, b) = weights(&loaded);
504        assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]);
505        assert_eq!(b, vec![5.0, 6.0]);
506    }
507
508    #[test]
509    fn records_exclude_flags_and_loading_preserves_the_target_flag() {
510        let device = Default::default();
511        let bytes = Flagged::new(7.0, false, &device)
512            .into_record()
513            .into_bytes()
514            .unwrap();
515        let record = ModuleRecord::from_bytes(bytes).unwrap();
516
517        assert_eq!(record.len(), 1, "only the tensor should be recorded");
518
519        let target = Flagged::new(0.0, true, &device);
520        let target_flag_id = target.enabled.id;
521        let loaded = target.load_record(record.clone());
522        assert_eq!(loaded.weight.val().try_into_vec_as::<f32>().unwrap(), [7.0]);
523        assert!(loaded.enabled.is_enabled());
524        assert_eq!(loaded.enabled.id, target_flag_id);
525
526        let target = Flagged::new(0.0, false, &device);
527        let target_flag_id = target.enabled.id;
528        let loaded = target.load_record(record);
529        assert!(!loaded.enabled.is_enabled());
530        assert_eq!(loaded.enabled.id, target_flag_id);
531    }
532
533    #[test]
534    fn a_group_records_its_own_parameters_only() {
535        let device = Default::default();
536        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
537
538        let record = model
539            .clone()
540            .into_record_group(ParamGroup::from_path("weight"));
541        assert_eq!(record.len(), 1);
542
543        // Applied back over a module the record says nothing about the rest of:
544        // the group's parameter lands, everything else keeps what it had.
545        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device)
546            .try_load_record(record.allow_partial(true))
547            .unwrap();
548        let (weight, bias) = weights(&loaded);
549        assert_eq!(weight, vec![1.0, 2.0, 3.0, 4.0]);
550        assert_eq!(bias, vec![0.0, 0.0], "the bias is outside the group");
551    }
552
553    #[test]
554    fn a_record_entry_matching_no_parameter_is_refused() {
555        let device = Default::default();
556        // A record of the wider module, applied to one that has no `gamma`: the entry lands
557        // nowhere, and a load that quietly did less than it was asked is what this refuses.
558        let record = TinyWide::zeros(&device).into_record();
559
560        let refusal = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).try_load_record(record.clone());
561        let Err(RecordError::Validation(message)) = refusal else {
562            panic!("a record entry that matches no parameter must be refused");
563        };
564        assert!(
565            message.contains("gamma"),
566            "the refusal must name it: {message}"
567        );
568
569        // Allowed for the deliberate case — loading a checkpoint into a part of the module it
570        // came from — and then the parameters that do match still land.
571        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device)
572            .try_load_record(record.allow_unused(true))
573            .unwrap();
574        assert_eq!(weights(&loaded), (vec![0.0; 4], vec![0.0; 2]));
575    }
576
577    #[test]
578    fn round_trip_file() {
579        let device = Default::default();
580        let dir = tempfile::tempdir().unwrap();
581        let path = dir.path().join("tiny.bpk");
582
583        Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device)
584            .into_record()
585            .save(&path)
586            .unwrap();
587
588        let record = ModuleRecord::load(&path).unwrap();
589        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
590        let (w, b) = weights(&loaded);
591        assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]);
592        assert_eq!(b, vec![5.0, 6.0]);
593    }
594
595    #[test]
596    fn missing_tensor_requires_allow_partial() {
597        let device = Default::default();
598        // Tiny has weight+bias; TinyWide also expects `gamma`, which the record lacks.
599        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
600
601        let strict = TinyWide::zeros(&device).try_load_record(record.clone());
602        assert!(matches!(strict, Err(RecordError::Validation(_))));
603
604        let partial = TinyWide::zeros(&device).try_load_record(record.allow_partial(true));
605        assert!(partial.is_ok());
606        let loaded = partial.unwrap();
607        // weight/bias were loaded; gamma kept its (zero) initialization.
608        assert_eq!(
609            loaded.weight.val().try_into_vec_as::<f32>().unwrap(),
610            vec![1.0, 2.0, 3.0, 4.0]
611        );
612        assert_eq!(
613            loaded.gamma.val().try_into_vec_as::<f32>().unwrap(),
614            vec![0.0, 0.0]
615        );
616    }
617
618    /// Build a `Tiny` whose parameters are zero-valued and carry the given dtype.
619    fn tiny_with_dtype(device: &Device, dtype: DType) -> Tiny {
620        Tiny {
621            weight: Param::from_tensor(
622                Tensor::<2>::from_data([[0.0, 0.0], [0.0, 0.0]], device).cast(dtype),
623            ),
624            bias: Param::from_tensor(Tensor::<1>::from_data([0.0, 0.0], device).cast(dtype)),
625        }
626    }
627
628    #[test]
629    fn dtype_policy_from_record_keeps_record_dtype() {
630        let device = Default::default();
631        // Record holds f32 data.
632        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
633
634        // Target params are f64, but the default policy (FromRecord) loads the data verbatim (f32).
635        let loaded = tiny_with_dtype(&device, DType::F64).load_record(record);
636        assert_eq!(loaded.weight.val().dtype(), DType::F32);
637        assert_eq!(loaded.bias.val().dtype(), DType::F32);
638    }
639
640    #[test]
641    fn dtype_policy_cast_to_module_uses_module_dtype() {
642        let device = Default::default();
643        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
644
645        // CastToModule casts the record's f32 data to the module's f64 params on load.
646        let loaded =
647            tiny_with_dtype(&device, DType::F64).load_record(record.cast_to_module_dtype());
648        assert_eq!(loaded.weight.val().dtype(), DType::F64);
649        assert_eq!(loaded.bias.val().dtype(), DType::F64);
650        // Values survive the cast.
651        assert_eq!(
652            loaded.weight.val().try_into_vec_as::<f64>().unwrap(),
653            vec![1.0, 2.0, 3.0, 4.0]
654        );
655    }
656
657    /// Build a `Tiny` whose `bias` shape (len 3) does not match the record's (len 2).
658    fn tiny_wrong_bias_shape(device: &Device) -> Tiny {
659        Tiny {
660            weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device),
661            bias: Param::from_data([0.0, 0.0, 0.0], device),
662        }
663    }
664
665    #[test]
666    fn shape_mismatch_fails_validation() {
667        let device = Default::default();
668        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
669
670        // `bias` shape mismatch is reported as a validation error by default.
671        let result = tiny_wrong_bias_shape(&device).try_load_record(record);
672        assert!(matches!(result, Err(RecordError::Validation(_))));
673    }
674
675    #[test]
676    fn load_record_preserves_param_id() {
677        let device = Default::default();
678        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
679
680        // Capture original ParamIds before saving.
681        let weight_id = model.weight.id;
682        let bias_id = model.bias.id;
683
684        let bytes = model.into_record().into_bytes().unwrap();
685        let record = ModuleRecord::from_bytes(bytes).unwrap();
686
687        // Load into a fresh model (new init() = different ParamIds).
688        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
689
690        // The loaded model should have the ORIGINAL ParamIds from the record,
691        // not the fresh ones from init().
692        assert_eq!(
693            loaded.weight.id, weight_id,
694            "weight ParamId should be restored from record"
695        );
696        assert_eq!(
697            loaded.bias.id, bias_id,
698            "bias ParamId should be restored from record"
699        );
700    }
701
702    #[test]
703    fn validate_false_ignores_shape_mismatch() {
704        let device = Default::default();
705        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
706
707        // With validation disabled, the shape-mismatched `bias` is skipped (keeps its value)
708        // while the matching `weight` still loads.
709        let loaded = tiny_wrong_bias_shape(&device)
710            .try_load_record(record.validate(false))
711            .unwrap();
712        assert_eq!(
713            loaded.weight.val().try_into_vec_as::<f32>().unwrap(),
714            vec![1.0, 2.0, 3.0, 4.0]
715        );
716        assert_eq!(
717            loaded.bias.val().try_into_vec_as::<f32>().unwrap(),
718            vec![0.0, 0.0, 0.0]
719        );
720    }
721
722    /// Mirrors a `Col`-layout `Linear` weight: persisted as `[3, 2]`, live as
723    /// the transposed `[2, 3]` through the init/save/load mappers.
724    #[derive(Module, Debug)]
725    struct ColLike {
726        weight: Param<Tensor<2>>,
727    }
728
729    impl ColLike {
730        fn new(seed: f32, device: &Device) -> Self {
731            let init_device = device.clone();
732            let weight = Param::uninitialized(
733                crate::module::ParamId::new(),
734                move |device, _| Tensor::<2>::full([3, 2], seed, device),
735                init_device,
736                true,
737                [3, 2].into(),
738            )
739            .init_mapper(|t: Tensor<2>| t.transpose())
740            .save_mapper(|t: Tensor<2>| t.transpose())
741            .load_mapper(|t: Tensor<2>| t.transpose());
742            Self { weight }
743        }
744    }
745
746    /// A param whose mapper changes the shape must round-trip through the
747    /// record in its save form.
748    #[test]
749    fn round_trip_a_shape_mapped_param() {
750        let device = Default::default();
751
752        let saved = ColLike::new(1.0, &device);
753        assert_eq!(saved.weight.val().dims(), [2, 3]);
754        let record = saved.into_record();
755        assert_eq!(
756            record.tensors[0].data.shape,
757            Shape::from([3, 2]),
758            "the record must hold the save form, not the live form"
759        );
760
761        let record = ModuleRecord::from_bytes(record.into_bytes().unwrap()).unwrap();
762        let loaded = ColLike::new(0.0, &device).load_record(record);
763        assert_eq!(loaded.weight.val().dims(), [2, 3]);
764        assert_eq!(
765            loaded.weight.val().try_into_vec_as::<f32>().unwrap(),
766            vec![1.0; 6],
767            "the recorded values must land, mapped back to the live form"
768        );
769    }
770}
burn_core/store/mod.rs

burn_core/store/
mod.rs
