burn_core/module/param/
reparameterization.rs

1use alloc::{string::String, string::ToString, vec::Vec};
2use burn_tensor::Tensor;
3
4use crate::module::{AutodiffModule, 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: AutodiffModule + 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::{module::Module, test_device, test_utils::SimpleLinear};
81    use burn_tensor::{Shape, Tolerance};
82
83    #[derive(Debug, Module)]
84    struct CustomScale {
85        scale: Param<Tensor<1>>,
86    }
87
88    impl Reparameterization for CustomScale {
89        const NAME: &'static str = "custom_scale";
90
91        fn materialize<const D: usize>(&self, base: Tensor<D>) -> Tensor<D> {
92            base * self.scale.val().reshape(Shape::from(alloc::vec![1; D]))
93        }
94    }
95
96    struct CustomScaleMapper;
97
98    impl Reparameterizer for CustomScaleMapper {
99        type Reparam = CustomScale;
100
101        fn reparameterize<const D: usize>(
102            &mut self,
103            _path: &str,
104            param: Param<Tensor<D>>,
105        ) -> (Param<Tensor<D>>, Option<Self::Reparam>) {
106            if D != 2 {
107                return (param, None);
108            }
109            let scale = Tensor::<1>::ones([1], &param.lazy_device());
110            (
111                param,
112                Some(CustomScale {
113                    scale: Param::from_tensor(scale),
114                }),
115            )
116        }
117    }
118
119    #[test]
120    fn custom_reparameterization_supports_full_module_lifecycle() {
121        let device = test_device().autodiff();
122        let model = SimpleLinear::new(4, 6, &device).apply_reparameterization(CustomScaleMapper);
123        let custom = model
124            .weight
125            .reparameterization::<CustomScale>()
126            .expect("custom reparameterization should be attached");
127
128        model
129            .weight
130            .val()
131            .into_data()
132            .assert_approx_eq::<f32>(&model.weight.base().into_data(), Tolerance::default());
133        assert_eq!(model.num_params(), 24 + 6 + 1);
134
135        let grads = model.weight.val().sum().backward();
136        assert!(model.weight.base().grad(&grads).is_some());
137        assert!(custom.scale.val().grad(&grads).is_some());
138
139        let target = SimpleLinear::new(4, 6, &device).apply_reparameterization(CustomScaleMapper);
140        let loaded = target.load_record(model.clone().into_record());
141        loaded
142            .weight
143            .val()
144            .into_data()
145            .assert_approx_eq::<f32>(&model.weight.val().into_data(), Tolerance::default());
146
147        let inference = model.valid();
148        assert!(inference.weight.reparameterization_dyn().is_none());
149    }
150}
burn_core/module/param/reparameterization.rs

burn_core/module/param/
reparameterization.rs
