minidx_core/layers/
diag.rs

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

minidx_core/layers/
diag.rs
