burn-store 0.22.1

Storage and serialization infrastructure for Burn
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use burn_core as burn;

use crate::{ModuleSnapshot, SafetensorsStore};
use burn_core::module::{Module, Param};
use burn_core::tensor::{Device, Tensor, shape};
use burn_nn::{Linear, LinearConfig};

#[derive(Module, Debug)]
pub(super) struct ComplexModule {
    pub encoder: EncoderModule,
    pub decoder: DecoderModule,
    pub layers: Vec<Linear>,
}

#[derive(Module, Debug)]
pub(super) struct EncoderModule {
    pub weight: Param<Tensor<3>>,
    pub bias: Param<Tensor<1>>,
    pub norm: Param<Tensor<1>>,
}

#[derive(Module, Debug)]
pub(super) struct DecoderModule {
    pub projection: Linear,
    pub scale: Param<Tensor<2>>,
}

impl ComplexModule {
    pub fn new(device: &Device) -> Self {
        Self {
            encoder: EncoderModule {
                weight: Param::from_data(
                    [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]],
                    device,
                ),
                bias: Param::from_data([0.1, 0.2, 0.3], device),
                norm: Param::from_data([1.0, 1.0, 1.0], device),
            },
            decoder: DecoderModule {
                projection: LinearConfig::new(4, 2).with_bias(true).init(device),
                scale: Param::from_data([[0.5, 0.5], [0.5, 0.5]], device),
            },
            layers: vec![
                LinearConfig::new(3, 4).with_bias(false).init(device),
                LinearConfig::new(4, 3).with_bias(true).init(device),
            ],
        }
    }

    pub fn new_zeros(device: &Device) -> Self {
        Self {
            encoder: EncoderModule {
                weight: Param::from_tensor(Tensor::zeros([2, 2, 2], device)),
                bias: Param::from_tensor(Tensor::zeros([3], device)),
                norm: Param::from_tensor(Tensor::zeros([3], device)),
            },
            decoder: DecoderModule {
                projection: LinearConfig::new(4, 2).with_bias(true).init(device),
                scale: Param::from_tensor(Tensor::zeros([2, 2], device)),
            },
            layers: vec![
                LinearConfig::new(3, 4).with_bias(false).init(device),
                LinearConfig::new(4, 3).with_bias(true).init(device),
            ],
        }
    }
}

#[test]
fn complex_module_round_trip() {
    let device = Default::default();
    let module1 = ComplexModule::new(&device);
    let mut module2 = ComplexModule::new_zeros(&device);

    // Save module1 using new store API
    let mut save_store = SafetensorsStore::from_bytes(None);
    module1.save_into(&mut save_store).unwrap();

    // Load into module2
    let mut load_store = SafetensorsStore::from_bytes(None);
    if let SafetensorsStore::Memory(ref mut p) = load_store
        && let SafetensorsStore::Memory(ref p_save) = save_store
    {
        // Get Arc and extract data
        let data_arc = p_save.data().unwrap();
        p.set_data(data_arc.as_ref().clone());
    }
    let result = module2.load_from(&mut load_store).unwrap();

    assert!(result.is_success());
    assert!(result.applied.len() > 5);
    assert_eq!(result.errors.len(), 0);

    // Verify data was imported correctly
    let module2_views = module2.collect(None, None, false);
    let encoder_weight = module2_views
        .iter()
        .find(|v| v.name == "encoder.weight")
        .map(|v| crate::bridge::to_data(v).unwrap())
        .unwrap();
    assert_eq!(*encoder_weight.shape(), shape![2, 2, 2]);
}

/// Safetensors stores no parameter identity, so importing from it must leave the target module's
/// ParamIds untouched. Randomizing them here would orphan optimizer state keyed by ParamId.
#[test]
fn round_trip_keeps_target_param_ids() {
    let device = Default::default();
    let module1 = ComplexModule::new(&device);
    let mut module2 = ComplexModule::new_zeros(&device);

    let expected_bias_id = module2.encoder.bias.id;
    let expected_norm_id = module2.encoder.norm.id;

    let mut save_store = SafetensorsStore::from_bytes(None);
    module1.save_into(&mut save_store).unwrap();

    let mut load_store = SafetensorsStore::from_bytes(None);
    if let SafetensorsStore::Memory(ref mut p) = load_store
        && let SafetensorsStore::Memory(ref p_save) = save_store
    {
        let data_arc = p_save.data().unwrap();
        p.set_data(data_arc.as_ref().clone());
    }
    let result = module2.load_from(&mut load_store).unwrap();
    assert!(result.is_success());

    assert_eq!(
        module2.encoder.bias.id, expected_bias_id,
        "encoder.bias should keep its ParamId across a safetensors import"
    );
    assert_eq!(
        module2.encoder.norm.id, expected_norm_id,
        "encoder.norm should keep its ParamId across a safetensors import"
    );
}