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

burn_core/store/
mod.rs
