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

//! End-to-end GPU integration test: XOR training with the entire
//! forward-backward-optimizer pipeline running on the GPU.
//!
//! Gated by `#[cfg(feature = "gpu")]` — skipped in CPU-only builds.
//! Also gracefully skips at runtime if no GPU adapter is available.

#![cfg(feature = "gpu")]

use rumus::autograd;
use rumus::backend::gpu::context::GpuContext;
use rumus::nn::{self, Linear, Module, ModuleToGpu};
use rumus::optim::{Adam, Optimizer};
use rumus::tensor::Tensor;

#[derive(Module)]
struct XorMLP {
    hidden: Linear,
    output: Linear,
}

impl XorMLP {
    fn new() -> Self {
        Self {
            hidden: Linear::new(2, 8, true),
            output: Linear::new(8, 1, true),
        }
    }

    fn forward(&self, x: &Tensor) -> Tensor {
        let h = nn::relu(&self.hidden.forward(x));
        self.output.forward(&h)
    }
}

#[test]
fn test_gpu_xor_training() {
    // Skip if no GPU available (CI-friendly).
    if !GpuContext::is_available() {
        eprintln!("Skipping GPU test: no compatible GPU adapter found");
        return;
    }

    // ---- XOR dataset ----
    let inputs = Tensor::new(
        vec![0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0],
        vec![4, 2],
    );
    let targets = Tensor::new(vec![0.0, 1.0, 1.0, 0.0], vec![4, 1]);

    // ---- Model ----
    let model = XorMLP::new();

    // Push everything to GPU.
    model.to_gpu();
    inputs.to_gpu();
    targets.to_gpu();

    let mut optimizer = Adam::new(model.parameters(), 0.01);

    // ---- Training loop ----
    let mut final_loss = f32::MAX;
    for _epoch in 0..200 {
        let pred = model.forward(&inputs);
        let loss = nn::mse_loss(&pred, &targets);

        // mse_loss returns a scalar — reading it triggers a 4-byte D2H.
        {
            let g = loss.data();
            final_loss = g[0];
        }

        let mut grads = autograd::backward(&loss).expect("backward failed");
        optimizer.step(&mut grads).expect("optimizer step failed");
    }

    // ---- Assert convergence ----
    assert!(
        final_loss < 0.01,
        "GPU XOR training did not converge: final loss = {:.6}",
        final_loss,
    );
}