minidx_core/layers/
bias1d.rs

1use crate::gradients::{ClassBias, ClassWrapper, Gradients};
2use crate::Dtype;
3
4/// A learnable bias on each element.
5#[derive(Clone, Debug)]
6pub struct Bias1d<E: Dtype, const I: usize> {
7    pub(crate) bias: ClassWrapper<[E; I], ClassBias>,
8}
9
10impl<E: Dtype, const I: usize> Default for Bias1d<E, I> {
11    fn default() -> Self {
12        Self {
13            bias: ClassWrapper::<[E; I], ClassBias>::wrap([E::default(); I]),
14        }
15    }
16}
17
18impl<E: Dtype, const I: usize> Bias1d<E, I> {
19    fn forward(&self, input: &[E; I]) -> [E; I] {
20        let mut out: [E; I] = self.bias.raw_grads_ref().clone();
21        for (o, i) in out.iter_mut().zip(input.iter()) {
22            *o += *i;
23        }
24        out
25    }
26}
27
28impl<E: Dtype, const I: usize> crate::BaseModule for Bias1d<E, I> {}
29
30impl<E: Dtype, const I: usize> crate::Module<[E; I]> for Bias1d<E, I> {
31    type Output = [E; I];
32
33    fn forward(&self, x: &[E; I]) -> Result<Self::Output, crate::Error> {
34        Ok(Bias1d::forward(self, x))
35    }
36}
37
38impl<E: Dtype, const I: usize> crate::RevModule<[E; I]> for Bias1d<E, I> {
39    type SelfGrads = ClassWrapper<[E; I], ClassBias>;
40
41    fn reverse(&self, _inputs: &[E; I], grads_wrt_output: &[E; I]) -> ([E; I], Self::SelfGrads) {
42        (
43            grads_wrt_output.clone(),
44            Self::SelfGrads::wrap(grads_wrt_output.clone()),
45        )
46    }
47
48    fn apply(
49        &mut self,
50        applyer: &mut impl crate::optimizers::GradApplyer,
51        updates: Self::SelfGrads,
52    ) -> Result<(), crate::Error> {
53        applyer.apply(updates, &mut self.bias)
54    }
55}
56
57impl<E: Dtype, const I: usize> crate::ResetParams for Bias1d<E, I> {
58    fn rand_params<RNG: rand::Rng>(
59        &mut self,
60        rng: &mut RNG,
61        scale: f32,
62    ) -> Result<(), crate::Error> {
63        // Xavier/Glorot initialization vibes, but scaled down a ton for bias initialization.
64        // (Unlike dense layers, biases are still learned from zero parameters)
65        let stddev = 1.0 / ((I * I) as f32 * 64.0).sqrt();
66        let normal = rand_distr::Normal::new(0.0, stddev).unwrap();
67
68        self.bias.grad_iter_mut().for_each(|b| {
69            let s: f32 = rng.sample::<f32, _>(normal) * scale;
70            *b = E::from_f32(s).unwrap();
71        });
72        Ok(())
73    }
74}
75
76impl<E: Dtype, const I: usize> crate::VisualizableUnit for Bias1d<E, I> {
77    const KIND: &'static str = "bias1d";
78    type Params = [[E; I]; 1];
79    fn params(&self) -> &Self::Params {
80        // SAFETY: An array of N is exactly the same as a unary array of the array of N
81        unsafe { std::mem::transmute(self.bias.raw_grads_ref()) }
82    }
83}
84
85impl<E: Dtype, const I: usize> crate::LoadableModule for Bias1d<E, I> {
86    fn save(
87        &self,
88        path: String,
89        dict: &mut std::collections::HashMap<String, Vec<f64>>,
90    ) -> Result<(), crate::LoadSaveError> {
91        dict.insert(
92            path,
93            self.bias.grad_iter().map(|f| f.to_f64().unwrap()).collect(),
94        );
95        Ok(())
96    }
97
98    fn load(
99        &mut self,
100        path: String,
101        dict: &std::collections::HashMap<String, Vec<f64>>,
102    ) -> Result<(), crate::LoadSaveError> {
103        let params = dict.get(&path).ok_or(crate::LoadSaveError {
104            path: path.clone(),
105            err: "Parameters missing".into(),
106        })?;
107        if params.len() != I {
108            return Err(crate::LoadSaveError {
109                path,
110                err: format!(
111                    "Parameters have wrong size: got {}, want {}",
112                    params.len(),
113                    I
114                )
115                .into(),
116            });
117        }
118        for (a, b) in self.bias.grad_iter_mut().zip(params.into_iter()) {
119            *a = E::from_f64(*b).unwrap();
120        }
121        Ok(())
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn test_zero() {
131        let layer = Bias1d::<f32, 1>::default();
132        assert_eq!(layer.forward(&[1.0]), [1.0],);
133    }
134
135    #[test]
136    fn test_2() {
137        let layer = Bias1d::<f32, 2> {
138            bias: ClassWrapper::<[f32; 2], ClassBias>::wrap([1.5, 3.2]),
139        };
140        assert_eq!(layer.forward(&[1.0, 3.0]), [2.5, 6.2],);
141    }
142}
minidx_core/layers/bias1d.rs

minidx_core/layers/
bias1d.rs
