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`](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    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                let id = t.param_id.map(ParamId::from).unwrap_or_else(ParamId::new);
210                let data = TensorData::from_bytes(t.bytes, t.shape, t.dtype);
211                RecordTensor {
212                    path: t.name,
213                    id,
214                    data,
215                }
216            })
217            .collect();
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::{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 TinyWide {
456        weight: Param<Tensor<2>>,
457        bias: Param<Tensor<1>>,
458        gamma: Param<Tensor<1>>,
459    }
460
461    impl TinyWide {
462        fn zeros(device: &Device) -> Self {
463            Self {
464                weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device),
465                bias: Param::from_data([0.0, 0.0], device),
466                gamma: Param::from_data([0.0, 0.0], device),
467            }
468        }
469    }
470
471    fn weights(model: &Tiny) -> (Vec<f32>, Vec<f32>) {
472        (
473            model.weight.val().to_data().to_vec().unwrap(),
474            model.bias.val().to_data().to_vec().unwrap(),
475        )
476    }
477
478    #[test]
479    fn round_trip_in_memory() {
480        let device = Default::default();
481        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
482
483        let bytes = model.into_record().into_bytes().unwrap();
484        let record = ModuleRecord::from_bytes(bytes).unwrap();
485        assert_eq!(record.len(), 2);
486
487        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
488        let (w, b) = weights(&loaded);
489        assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]);
490        assert_eq!(b, vec![5.0, 6.0]);
491    }
492
493    #[test]
494    fn a_group_records_its_own_parameters_only() {
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 record = model
499            .clone()
500            .into_record_group(ParamGroup::from_path("weight"));
501        assert_eq!(record.len(), 1);
502
503        // Applied back over a module the record says nothing about the rest of:
504        // the group's parameter lands, everything else keeps what it had.
505        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device)
506            .try_load_record(record.allow_partial(true))
507            .unwrap();
508        let (weight, bias) = weights(&loaded);
509        assert_eq!(weight, vec![1.0, 2.0, 3.0, 4.0]);
510        assert_eq!(bias, vec![0.0, 0.0], "the bias is outside the group");
511    }
512
513    #[test]
514    fn a_record_entry_matching_no_parameter_is_refused() {
515        let device = Default::default();
516        // A record of the wider module, applied to one that has no `gamma`: the entry lands
517        // nowhere, and a load that quietly did less than it was asked is what this refuses.
518        let record = TinyWide::zeros(&device).into_record();
519
520        let refusal = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).try_load_record(record.clone());
521        let Err(RecordError::Validation(message)) = refusal else {
522            panic!("a record entry that matches no parameter must be refused");
523        };
524        assert!(
525            message.contains("gamma"),
526            "the refusal must name it: {message}"
527        );
528
529        // Allowed for the deliberate case — loading a checkpoint into a part of the module it
530        // came from — and then the parameters that do match still land.
531        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device)
532            .try_load_record(record.allow_unused(true))
533            .unwrap();
534        assert_eq!(weights(&loaded), (vec![0.0; 4], vec![0.0; 2]));
535    }
536
537    #[test]
538    fn round_trip_file() {
539        let device = Default::default();
540        let dir = tempfile::tempdir().unwrap();
541        let path = dir.path().join("tiny.bpk");
542
543        Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device)
544            .into_record()
545            .save(&path)
546            .unwrap();
547
548        let record = ModuleRecord::load(&path).unwrap();
549        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
550        let (w, b) = weights(&loaded);
551        assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]);
552        assert_eq!(b, vec![5.0, 6.0]);
553    }
554
555    #[test]
556    fn missing_tensor_requires_allow_partial() {
557        let device = Default::default();
558        // Tiny has weight+bias; TinyWide also expects `gamma`, which the record lacks.
559        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
560
561        let strict = TinyWide::zeros(&device).try_load_record(record.clone());
562        assert!(matches!(strict, Err(RecordError::Validation(_))));
563
564        let partial = TinyWide::zeros(&device).try_load_record(record.allow_partial(true));
565        assert!(partial.is_ok());
566        let loaded = partial.unwrap();
567        // weight/bias were loaded; gamma kept its (zero) initialization.
568        assert_eq!(
569            loaded.weight.val().to_data().to_vec::<f32>().unwrap(),
570            vec![1.0, 2.0, 3.0, 4.0]
571        );
572        assert_eq!(
573            loaded.gamma.val().to_data().to_vec::<f32>().unwrap(),
574            vec![0.0, 0.0]
575        );
576    }
577
578    /// Build a `Tiny` whose parameters are zero-valued and carry the given dtype.
579    fn tiny_with_dtype(device: &Device, dtype: DType) -> Tiny {
580        Tiny {
581            weight: Param::from_tensor(
582                Tensor::<2>::from_data([[0.0, 0.0], [0.0, 0.0]], device).cast(dtype),
583            ),
584            bias: Param::from_tensor(Tensor::<1>::from_data([0.0, 0.0], device).cast(dtype)),
585        }
586    }
587
588    #[test]
589    fn dtype_policy_from_record_keeps_record_dtype() {
590        let device = Default::default();
591        // Record holds f32 data.
592        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
593
594        // Target params are f64, but the default policy (FromRecord) loads the data verbatim (f32).
595        let loaded = tiny_with_dtype(&device, DType::F64).load_record(record);
596        assert_eq!(loaded.weight.val().dtype(), DType::F32);
597        assert_eq!(loaded.bias.val().dtype(), DType::F32);
598    }
599
600    #[test]
601    fn dtype_policy_cast_to_module_uses_module_dtype() {
602        let device = Default::default();
603        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
604
605        // CastToModule casts the record's f32 data to the module's f64 params on load.
606        let loaded =
607            tiny_with_dtype(&device, DType::F64).load_record(record.cast_to_module_dtype());
608        assert_eq!(loaded.weight.val().dtype(), DType::F64);
609        assert_eq!(loaded.bias.val().dtype(), DType::F64);
610        // Values survive the cast.
611        assert_eq!(
612            loaded.weight.val().to_data().to_vec::<f64>().unwrap(),
613            vec![1.0, 2.0, 3.0, 4.0]
614        );
615    }
616
617    /// Build a `Tiny` whose `bias` shape (len 3) does not match the record's (len 2).
618    fn tiny_wrong_bias_shape(device: &Device) -> Tiny {
619        Tiny {
620            weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device),
621            bias: Param::from_data([0.0, 0.0, 0.0], device),
622        }
623    }
624
625    #[test]
626    fn shape_mismatch_fails_validation() {
627        let device = Default::default();
628        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
629
630        // `bias` shape mismatch is reported as a validation error by default.
631        let result = tiny_wrong_bias_shape(&device).try_load_record(record);
632        assert!(matches!(result, Err(RecordError::Validation(_))));
633    }
634
635    #[test]
636    fn load_record_preserves_param_id() {
637        let device = Default::default();
638        let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device);
639
640        // Capture original ParamIds before saving.
641        let weight_id = model.weight.id;
642        let bias_id = model.bias.id;
643
644        let bytes = model.into_record().into_bytes().unwrap();
645        let record = ModuleRecord::from_bytes(bytes).unwrap();
646
647        // Load into a fresh model (new init() = different ParamIds).
648        let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record);
649
650        // The loaded model should have the ORIGINAL ParamIds from the record,
651        // not the fresh ones from init().
652        assert_eq!(
653            loaded.weight.id, weight_id,
654            "weight ParamId should be restored from record"
655        );
656        assert_eq!(
657            loaded.bias.id, bias_id,
658            "bias ParamId should be restored from record"
659        );
660    }
661
662    #[test]
663    fn validate_false_ignores_shape_mismatch() {
664        let device = Default::default();
665        let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record();
666
667        // With validation disabled, the shape-mismatched `bias` is skipped (keeps its value)
668        // while the matching `weight` still loads.
669        let loaded = tiny_wrong_bias_shape(&device)
670            .try_load_record(record.validate(false))
671            .unwrap();
672        assert_eq!(
673            loaded.weight.val().to_data().to_vec::<f32>().unwrap(),
674            vec![1.0, 2.0, 3.0, 4.0]
675        );
676        assert_eq!(
677            loaded.bias.val().to_data().to_vec::<f32>().unwrap(),
678            vec![0.0, 0.0, 0.0]
679        );
680    }
681
682    /// Mirrors a `Col`-layout `Linear` weight: persisted as `[3, 2]`, live as
683    /// the transposed `[2, 3]` through the init/save/load mappers.
684    #[derive(Module, Debug)]
685    struct ColLike {
686        weight: Param<Tensor<2>>,
687    }
688
689    impl ColLike {
690        fn new(seed: f32, device: &Device) -> Self {
691            let init_device = device.clone();
692            let weight = Param::uninitialized(
693                crate::module::ParamId::new(),
694                move |device, _| Tensor::<2>::full([3, 2], seed, device),
695                init_device,
696                true,
697                [3, 2].into(),
698            )
699            .init_mapper(|t: Tensor<2>| t.transpose())
700            .save_mapper(|t: Tensor<2>| t.transpose())
701            .load_mapper(|t: Tensor<2>| t.transpose());
702            Self { weight }
703        }
704    }
705
706    /// A param whose mapper changes the shape must round-trip through the
707    /// record in its save form.
708    #[test]
709    fn round_trip_a_shape_mapped_param() {
710        let device = Default::default();
711
712        let saved = ColLike::new(1.0, &device);
713        assert_eq!(saved.weight.val().dims(), [2, 3]);
714        let record = saved.into_record();
715        assert_eq!(
716            record.tensors[0].data.shape,
717            Shape::from([3, 2]),
718            "the record must hold the save form, not the live form"
719        );
720
721        let record = ModuleRecord::from_bytes(record.into_bytes().unwrap()).unwrap();
722        let loaded = ColLike::new(0.0, &device).load_record(record);
723        assert_eq!(loaded.weight.val().dims(), [2, 3]);
724        assert_eq!(
725            loaded.weight.val().to_data().to_vec::<f32>().unwrap(),
726            vec![1.0; 6],
727            "the recorded values must land, mapped back to the live form"
728        );
729    }
730}
burn_core/store/mod.rs

burn_core/store/
mod.rs
