hodu_nn 0.1.2

hodu nn
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

use crate::module::Module;
use hodu_core::{error::HoduResult, scalar::Scalar, tensor::Tensor, types::dtype::DType};

#[derive(Module, Clone)]
pub struct Linear {
    weight: Tensor,
    bias: Option<Tensor>,
}

impl Linear {
    pub fn new(in_features: usize, out_features: usize, with_bias: bool, dtype: DType) -> HoduResult<Self> {
        let k: f32 = 1.0 / (in_features as f32).sqrt();

        let zero = Scalar::zero(dtype);
        let one = Scalar::one(dtype);
        let k_scalar = Scalar::from_f32(k, dtype);

        // weight: [out_features, in_features]
        let weight = Tensor::randn(&[out_features, in_features], zero, one)?;
        weight.set_requires_grad(true)?;
        // Scale by k (Xavier/Glorot initialization)
        let weight = weight.mul(&Tensor::full(&[], k_scalar)?)?;

        // bias
        let bias = if with_bias {
            let bias = Tensor::randn(&[out_features], zero, one)?;
            bias.set_requires_grad(true)?;
            // Scale by k
            let bias = bias.mul(&Tensor::full(&[], k_scalar)?)?;
            Some(bias)
        } else {
            None
        };

        Ok(Self { weight, bias })
    }

    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
        // Linear transformation: input @ weight.T + bias
        let output = input.matmul(&self.weight.transpose(-2, -1)?)?;

        if let Some(ref bias) = self.bias {
            output.add(bias)
        } else {
            Ok(output)
        }
    }

    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
        let mut params = vec![&mut self.weight];
        if let Some(ref mut bias) = self.bias {
            params.push(bias);
        }
        params
    }
}