rumus 0.3.1

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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! 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 {
        // ONNX tracing: emit a fused Gemm node instead of matmul + add_bias.
        #[cfg(feature = "onnx")]
        if crate::onnx::tracer::is_tracing() {
            crate::onnx::tracer::with_tracer(|t| t.enter_fusion());
        }

        let out = input.matmul(&self.weight.tensor);
        let result = match &self.bias {
            Some(bias) => out.add_bias(&bias.tensor),
            None => out,
        };

        #[cfg(feature = "onnx")]
        if crate::onnx::tracer::is_tracing() {
            crate::onnx::tracer::with_tracer(|t| t.leave_fusion());
            crate::onnx::tracer::trace_linear(
                input,
                &self.weight.tensor,
                self.bias.as_ref().map(|b| &b.tensor),
                &result,
                "weight",
                self.bias.as_ref().map(|_| "bias"),
            );
        }

        result
    }
}

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(())
    }
}