use rand::RngCore;
#[derive(Clone)]
pub struct Linear {
pub weights: Vec<f32>, pub biases: Vec<f32>,
input_size: usize,
output_size: usize,
}
impl Linear {
pub fn new(input_size: usize, output_size: usize, rng: &mut dyn RngCore) -> Self {
let scale = (2.0 / input_size as f32).sqrt();
let mut weights = vec![0.0; input_size * output_size];
let biases = vec![0.0; output_size];
for w in weights.iter_mut() {
*w = (rng.next_u32() as f32 / u32::MAX as f32 * 2.0 - 1.0) * scale;
}
Self {
weights,
biases,
input_size,
output_size,
}
}
pub fn forward(&self, input: &[f32]) -> Vec<f32> {
let mut out = vec![0.0; self.output_size];
for j in 0..self.output_size {
let mut sum = self.biases[j];
for i in 0..self.input_size {
sum += input[i] * self.weights[j * self.input_size + i];
}
out[j] = sum;
}
out
}
pub fn backward(
&self,
input: &[f32],
grad_output: &[f32],
) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
let mut grad_input = vec![0.0; self.input_size];
let mut grad_weights = vec![0.0; self.weights.len()];
let mut grad_biases = vec![0.0; self.output_size];
for j in 0..self.output_size {
let go = grad_output[j];
grad_biases[j] = go;
for i in 0..self.input_size {
grad_weights[j * self.input_size + i] = go * input[i];
grad_input[i] += go * self.weights[j * self.input_size + i];
}
}
(grad_input, grad_weights, grad_biases)
}
}
#[derive(Clone)]
pub struct MLP {
pub layers: Vec<Linear>,
pre_activations: Vec<Vec<f32>>,
inputs: Vec<Vec<f32>>,
}
impl MLP {
pub fn new(sizes: &[usize], rng: &mut dyn RngCore) -> Self {
let mut layers = Vec::new();
for i in 0..sizes.len() - 1 {
layers.push(Linear::new(sizes[i], sizes[i + 1], rng));
}
Self {
layers,
pre_activations: Vec::new(),
inputs: Vec::new(),
}
}
pub fn forward(&mut self, input: &[f32]) -> Vec<f32> {
self.pre_activations.clear();
self.inputs.clear();
let mut current = input.to_vec();
self.inputs.push(current.clone());
for layer in &self.layers {
let pre = layer.forward(¤t);
self.pre_activations.push(pre.clone());
if self.layers.len() > 1 && self.pre_activations.len() < self.layers.len() {
current = pre.iter().map(|&x| x.max(0.0)).collect();
} else {
current = pre;
}
self.inputs.push(current.clone());
}
current
}
pub fn backward(&mut self, grad_output: &[f32]) -> Vec<(Vec<f32>, Vec<f32>)> {
let mut grad = grad_output.to_vec();
let mut grads = Vec::new();
for i in (0..self.layers.len()).rev() {
if i < self.layers.len() - 1 {
let pre = &self.pre_activations[i];
for j in 0..grad.len() {
if pre[j] <= 0.0 {
grad[j] = 0.0;
}
}
}
let (grad_input, gw, gb) = self.layers[i].backward(&self.inputs[i], &grad);
grads.push((gw, gb));
grad = grad_input;
}
grads.reverse();
grads
}
pub fn update(&mut self, grads: &[(Vec<f32>, Vec<f32>)], lr: f32) {
for (i, (gw, gb)) in grads.iter().enumerate() {
for j in 0..self.layers[i].weights.len() {
self.layers[i].weights[j] -= lr * gw[j];
}
for j in 0..self.layers[i].biases.len() {
self.layers[i].biases[j] -= lr * gb[j];
}
}
}
}
pub fn softmax(x: &[f32]) -> Vec<f32> {
let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps: Vec<f32> = x.iter().map(|&v| (v - max).exp()).collect();
let sum: f32 = exps.iter().sum();
exps.iter().map(|&v| v / sum).collect()
}
pub fn log_softmax(x: &[f32]) -> Vec<f32> {
let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps: Vec<f32> = x.iter().map(|&v| (v - max).exp()).collect();
let sum: f32 = exps.iter().sum();
let log_sum = sum.ln();
x.iter().map(|&v| v - max - log_sum).collect()
}
pub fn cross_entropy_grad(logits: &[f32], target: usize) -> Vec<f32> {
let mut probs = softmax(logits);
probs[target] -= 1.0;
probs
}
pub fn mse_grad(pred: &[f32], target: &[f32]) -> Vec<f32> {
pred.iter()
.zip(target.iter())
.map(|(p, t)| 2.0 * (p - t))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use rand::rngs::StdRng;
use rand::SeedableRng;
#[test]
fn test_linear_forward_backward() {
let mut rng = StdRng::seed_from_u64(0);
let layer = Linear::new(3, 2, &mut rng);
let input = vec![1.0, 2.0, 3.0];
let out = layer.forward(&input);
assert_eq!(out.len(), 2);
let grad = vec![0.5, -0.5];
let (g_in, g_w, g_b) = layer.backward(&input, &grad);
assert_eq!(g_in.len(), 3);
assert_eq!(g_w.len(), 6);
assert_eq!(g_b.len(), 2);
}
#[test]
fn test_mlp_forward_backward() {
let mut rng = StdRng::seed_from_u64(0);
let mut mlp = MLP::new(&[2, 3, 2], &mut rng);
let out = mlp.forward(&[1.0, -1.0]);
assert_eq!(out.len(), 2);
let grads = mlp.backward(&[1.0, 0.0]);
assert_eq!(grads.len(), 2);
}
#[test]
fn test_softmax() {
let p = softmax(&[1.0, 2.0, 3.0]);
let sum: f32 = p.iter().sum();
assert!((sum - 1.0).abs() < 1e-5);
}
#[test]
fn test_log_softmax() {
let p = log_softmax(&[1.0, 2.0, 3.0]);
let exp_sum: f32 = p.iter().map(|&v| v.exp()).sum();
assert!((exp_sum - 1.0).abs() < 1e-5);
}
#[test]
fn test_cross_entropy_grad() {
let logits = vec![0.0, 1.0, 0.0];
let grad = cross_entropy_grad(&logits, 1);
assert!(grad[1] < 0.0);
}
#[test]
fn test_mse_grad() {
let pred = vec![1.0, 2.0, 3.0];
let target = vec![2.0, 2.0, 2.0];
let grad = mse_grad(&pred, &target);
assert_eq!(grad, vec![-2.0, 0.0, 2.0]);
}
#[test]
fn test_single_layer_mlp() {
let mut rng = StdRng::seed_from_u64(0);
let mut mlp = MLP::new(&[2, 3], &mut rng);
let out = mlp.forward(&[1.0, 0.5]);
assert_eq!(out.len(), 3);
let grads = mlp.backward(&[1.0, 0.5, -0.5]);
assert_eq!(grads.len(), 1);
}
#[test]
fn test_mlp_clone_independence() {
let mut rng = StdRng::seed_from_u64(0);
let mut mlp1 = MLP::new(&[2, 3], &mut rng);
let mut mlp2 = mlp1.clone();
mlp1.layers[0].weights[0] = 999.0;
assert_ne!(mlp1.layers[0].weights[0], mlp2.layers[0].weights[0]);
}
}