burn_core/module/param/
reparameterization.rs

1use alloc::{string::String, string::ToString, vec::Vec};
2use burn_tensor::Tensor;
3
4use crate::module::{Module, ModuleMapper};
5
6use super::Param;
7
8/// A rank-specific parameter reparameterization.
9///
10/// Implementations are regular [`Module`](crate::module::Module)s, so their parameters automatically participate in
11/// optimizer, record, device and autodiff traversal. The implementation only needs to describe
12/// how its state materializes an effective value from the stored base parameter.
13pub trait Reparameterization: Module + Sync + 'static {
14    /// Stable path component used for the reparameterization's nested parameters.
15    const NAME: &'static str;
16    /// Materialize the effective parameter value from its stored base.
17    fn materialize<const D: usize>(&self, base: Tensor<D>) -> Tensor<D>;
18}
19
20/// Defines how floating-point parameters are prepared for reparameterization.
21///
22/// [`Module::apply_reparameterization`](crate::module::Module::apply_reparameterization) passes
23/// every floating-point parameter encountered during module traversal to [`reparameterize`](Self::reparameterize).
24/// Implementations may use the parameter path to decide whether to attach a
25/// [`Reparameterization`] and may transform the parameter into the structural base that should be
26/// stored.
27pub trait Reparameterizer {
28    /// Reparameterization produced for a parameter.
29    type Reparam: Reparameterization;
30
31    /// Prepare a parameter and optionally create a reparameterization for it.
32    ///
33    /// The returned parameter is always used as the structural base. Returning `None` leaves that
34    /// base without a reparameterization.
35    fn reparameterize<const D: usize>(
36        &mut self,
37        path: &str,
38        param: Param<Tensor<D>>,
39    ) -> (Param<Tensor<D>>, Option<Self::Reparam>);
40}
41
42pub(crate) struct ApplyReparameterization<R> {
43    reparameterizer: R,
44    path: Vec<String>,
45}
46
47impl<R> ApplyReparameterization<R> {
48    pub(crate) fn new(reparameterizer: R) -> Self {
49        Self {
50            reparameterizer,
51            path: Vec::new(),
52        }
53    }
54}
55
56impl<R: Reparameterizer> ModuleMapper for ApplyReparameterization<R> {
57    fn enter_module(&mut self, name: &str, _container_type: &str) {
58        self.path.push(name.to_string());
59    }
60
61    fn exit_module(&mut self, _name: &str, _container_type: &str) {
62        self.path.pop();
63    }
64
65    fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
66        let path = self.path.join(".");
67        let (base, reparameterization) = self.reparameterizer.reparameterize(&path, param);
68        match reparameterization {
69            Some(reparameterization) => base.with_reparameterization(reparameterization),
70            None => base,
71        }
72    }
73}
74
75#[cfg(all(test, feature = "autodiff"))]
76mod tests {
77    use super::*;
78    use crate as burn;
79    use crate::module::Reparameterizer;
80    use crate::{
81        module::{Flag, Module, ParamGroup},
82        test_device,
83        test_utils::SimpleLinear,
84    };
85    use burn_tensor::{Shape, Tolerance};
86
87    #[derive(Debug, Module)]
88    struct CustomScale {
89        scale: Param<Tensor<1>>,
90        enabled: Param<Flag>,
91    }
92
93    impl Reparameterization for CustomScale {
94        const NAME: &'static str = "custom_scale";
95
96        fn materialize<const D: usize>(&self, base: Tensor<D>) -> Tensor<D> {
97            base * self.scale.val().reshape(Shape::from(alloc::vec![1; D]))
98        }
99    }
100
101    struct CustomScaleMapper;
102
103    impl Reparameterizer for CustomScaleMapper {
104        type Reparam = CustomScale;
105
106        fn reparameterize<const D: usize>(
107            &mut self,
108            _path: &str,
109            param: Param<Tensor<D>>,
110        ) -> (Param<Tensor<D>>, Option<Self::Reparam>) {
111            if D != 2 {
112                return (param, None);
113            }
114            let scale = Tensor::<1>::ones([1], &param.lazy_device());
115            (
116                param,
117                Some(CustomScale {
118                    scale: Param::from_tensor(scale),
119                    enabled: Param::from_bool(true),
120                }),
121            )
122        }
123    }
124
125    #[test]
126    fn custom_reparameterization_supports_full_module_lifecycle() {
127        let device = test_device().autodiff();
128        let model = SimpleLinear::new(4, 6, &device).apply_reparameterization(CustomScaleMapper);
129        let custom = model
130            .weight
131            .reparameterization::<CustomScale>()
132            .expect("custom reparameterization should be attached");
133
134        model
135            .weight
136            .val()
137            .into_data()
138            .assert_approx_eq::<f32>(&model.weight.base().into_data(), Tolerance::default());
139        assert_eq!(model.num_params(), 24 + 6 + 1);
140
141        let group = ParamGroup::ids_from_module(model.clone());
142        assert!(group.matches(&custom.enabled.id, None));
143
144        let frozen = model.clone().freeze();
145        let frozen_custom = frozen
146            .weight
147            .reparameterization::<CustomScale>()
148            .expect("custom reparameterization should remain attached");
149        assert!(!frozen_custom.enabled.is_enabled());
150
151        let grads = model.weight.val().sum().backward();
152        assert!(model.weight.base().grad(&grads).is_some());
153        assert!(custom.scale.val().grad(&grads).is_some());
154
155        let target = SimpleLinear::new(4, 6, &device).apply_reparameterization(CustomScaleMapper);
156        let loaded = target.load_record(model.clone().into_record());
157        loaded
158            .weight
159            .val()
160            .into_data()
161            .assert_approx_eq::<f32>(&model.weight.val().into_data(), Tolerance::default());
162
163        let inference = model.valid();
164        assert!(inference.weight.reparameterization_dyn().is_none());
165    }
166}
burn_core/module/param/reparameterization.rs

burn_core/module/param/
reparameterization.rs
