rumus 0.1.0

A native-Rust deep learning framework with explicit memory safety and hardware acceleration
Documentation 
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! Fully connected (dense) linear layer.

use std::cell::Cell;
use std::collections::HashMap;

use crate::autograd::AutogradError;
use crate::nn::{Module, Parameter};
use crate::tensor::Tensor;

// ---------------------------------------------------------------------------
// Simple LCG PRNG — zero external dependencies
// ---------------------------------------------------------------------------

thread_local! {
    static RNG_STATE: Cell<u64> = Cell::new(42);
}

/// Generate a pseudo-random f32 in `[-bound, +bound]` using a thread-local
/// LCG (linear congruential generator).
fn lcg_uniform(bound: f32) -> f32 {
    RNG_STATE.with(|state| {
        let s = state
            .get()
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        state.set(s);
        let u = (s >> 33) as f32 / (1u64 << 31) as f32;
        (2.0 * u - 1.0) * bound
    })
}

// ---------------------------------------------------------------------------
// Linear layer
// ---------------------------------------------------------------------------

/// Fully connected layer: `y = x @ weight + bias`.
///
/// Weight layout is `[in_features, out_features]` — avoids a transpose in
/// the forward pass, keeping the hot path allocation-free.
///
/// Initialization uses Kaiming Uniform:
/// `bound = sqrt(1 / in_features)`, each element ~ U(-bound, +bound).
pub struct Linear {
    /// Weight matrix, shape `[in_features, out_features]`.
    pub weight: Parameter,
    /// Optional bias vector, shape `[out_features]`.
    pub bias: Option<Parameter>,
}

impl Linear {
    /// Create a new `Linear` layer.
    pub fn new(in_features: usize, out_features: usize, with_bias: bool) -> Self {
        let bound = (1.0 / in_features as f32).sqrt();

        let weight_data: Vec<f32> = (0..in_features * out_features)
            .map(|_| lcg_uniform(bound))
            .collect();
        let weight =
            Parameter::new(Tensor::new(weight_data, vec![in_features, out_features]));

        let bias = if with_bias {
            let bias_data: Vec<f32> = (0..out_features).map(|_| lcg_uniform(bound)).collect();
            Some(Parameter::new(Tensor::new(bias_data, vec![out_features])))
        } else {
            None
        };

        Self { weight, bias }
    }

    /// Forward pass: `y = input @ weight + bias`.
    ///
    /// `input` shape: `[batch, in_features]`.
    /// Output shape: `[batch, out_features]`.
    pub fn forward(&self, input: &Tensor) -> Tensor {
        let out = input.matmul(&self.weight.tensor);
        match &self.bias {
            Some(bias) => out.add_bias(&bias.tensor),
            None => out,
        }
    }
}

impl Module for Linear {
    fn parameters(&self) -> Vec<Parameter> {
        let mut params = vec![self.weight.clone()];
        if let Some(ref bias) = self.bias {
            params.push(bias.clone());
        }
        params
    }

    fn state_dict(&self, prefix: &str) -> HashMap<String, Tensor> {
        let mut dict = self.weight.state_dict(&format!("{}weight.", prefix));
        if let Some(ref bias) = self.bias {
            dict.extend(bias.state_dict(&format!("{}bias.", prefix)));
        }
        dict
    }

    fn load_state_dict(
        &mut self,
        dict: &HashMap<String, Tensor>,
        prefix: &str,
    ) -> Result<(), AutogradError> {
        self.weight
            .load_state_dict(dict, &format!("{}weight.", prefix))?;
        if let Some(ref mut bias) = self.bias {
            bias.load_state_dict(dict, &format!("{}bias.", prefix))?;
        }
        Ok(())
    }
}