minidx_core/layers/
scalar_scale.rs

1use crate::Dtype;
2
3/// A layer which learns a scalar scaling factor.
4#[derive(Clone, Debug)]
5pub struct ScalarScale<E: Dtype> {
6    pub(crate) scale: E,
7}
8
9impl<E: Dtype> Default for ScalarScale<E> {
10    fn default() -> Self {
11        ScalarScale {
12            scale: E::ONE.div(E::ONE + E::ONE),
13        }
14    }
15}
16
17impl<E: Dtype> ScalarScale<E> {
18    #[inline]
19    fn forward<const I: usize>(&self, input: &[E; I]) -> [E; I] {
20        let mut out: [E; I] = [E::default(); I];
21
22        out.iter_mut()
23            .zip(input.iter())
24            .for_each(|(o, i)| *o = *i * self.scale);
25
26        out
27    }
28
29    #[inline]
30    fn gradients_wrt_input<const I: usize>(&self, output_gradients: &[E; I]) -> [E; I] {
31        let mut out: [E; I] = [E::default(); I];
32
33        out.iter_mut()
34            .zip(output_gradients.iter())
35            .for_each(|(o, w)| *o = self.scale * *w);
36
37        out
38    }
39
40    #[inline]
41    fn gradients_wrt_weights<const I: usize>(
42        &self,
43        input: &[E; I],
44        output_gradients: &[E; I],
45    ) -> E {
46        input
47            .iter()
48            .zip(output_gradients.iter())
49            .fold(E::default(), |a, (i, w)| a + *i * *w)
50    }
51}
52
53impl<E: Dtype> crate::BaseModule for ScalarScale<E> {}
54
55impl<E: Dtype, const I: usize> crate::Module<[E; I]> for ScalarScale<E> {
56    type Output = [E; I];
57
58    fn forward(&self, x: &[E; I]) -> Result<Self::Output, crate::Error> {
59        Ok(ScalarScale::forward(self, &x))
60    }
61}
62
63impl<E: Dtype, const I: usize> crate::RevModule<[E; I]> for ScalarScale<E> {
64    type SelfGrads = [E; 1];
65
66    fn reverse(&self, inputs: &[E; I], grads_wrt_output: &[E; I]) -> ([E; I], Self::SelfGrads) {
67        (
68            ScalarScale::gradients_wrt_input(self, grads_wrt_output),
69            [ScalarScale::gradients_wrt_weights(
70                self,
71                inputs,
72                grads_wrt_output,
73            )],
74        )
75    }
76
77    fn apply(
78        &mut self,
79        applyer: &mut impl crate::optimizers::GradApplyer,
80        updates: Self::SelfGrads,
81    ) -> Result<(), crate::Error> {
82        applyer.apply(updates, unsafe {
83            // SAFETY: A value of N has the same memory layout as stacked unary arrays of N
84            std::mem::transmute(&mut self.scale)
85        })
86    }
87}
88
89impl<E: Dtype> crate::ResetParams for ScalarScale<E> {
90    fn rand_params<RNG: rand::Rng>(
91        &mut self,
92        _rng: &mut RNG,
93        _scale: f32,
94    ) -> Result<(), crate::Error> {
95        // 0.5 initial value
96        self.scale = E::ONE.div(E::ONE + E::ONE);
97        Ok(())
98    }
99}
100
101impl<E: Dtype> crate::VisualizableUnit for ScalarScale<E> {
102    const KIND: &'static str = "scalar_scale";
103    type Params = [[E; 1]; 1];
104    fn params(&self) -> &Self::Params {
105        // SAFETY: A value of N has the same memory layout as stacked unary arrays of N
106        unsafe { std::mem::transmute(&self.scale) }
107    }
108}
109
110impl<E: Dtype> crate::LoadableModule for ScalarScale<E> {
111    fn save(
112        &self,
113        path: String,
114        dict: &mut std::collections::HashMap<String, Vec<f64>>,
115    ) -> Result<(), crate::LoadSaveError> {
116        dict.insert(path, vec![self.scale.to_f64().unwrap()]);
117        Ok(())
118    }
119
120    fn load(
121        &mut self,
122        path: String,
123        dict: &std::collections::HashMap<String, Vec<f64>>,
124    ) -> Result<(), crate::LoadSaveError> {
125        let params = dict.get(&path).ok_or(crate::LoadSaveError {
126            path: path.clone(),
127            err: "Parameters missing".into(),
128        })?;
129        if params.len() != 1 {
130            return Err(crate::LoadSaveError {
131                path,
132                err: format!(
133                    "Parameters have wrong size: got {}, want {}",
134                    params.len(),
135                    1
136                )
137                .into(),
138            });
139        }
140        self.scale = E::from_f64(params[0]).unwrap();
141        Ok(())
142    }
143}
minidx_core/layers/scalar_scale.rs

minidx_core/layers/
scalar_scale.rs
