axonml 0.6.2

AxonML — a complete ML/AI framework in pure Rust (umbrella crate)
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
136
137
138
139

//! Simple Training Example — XOR Problem with a Two-Layer MLP
//!
//! Minimal end-to-end training demonstration using the AxonML framework.
//! Prints the framework version and enabled features, constructs a two-layer
//! MLP (`Linear(2,4)` -> sigmoid -> `Linear(4,1)` -> sigmoid) to learn the
//! XOR function, trains for 1000 epochs with `Adam` optimizer (lr=0.1) and
//! manual MSE loss (`(output - target)^2`), logs loss every 200 epochs, and
//! evaluates all four XOR inputs with rounded prediction vs. expected output.
//!
//! # File
//! `crates/axonml/examples/simple_training.rs`
//!
//! # Author
//! Andrew Jewell Sr. — AutomataNexus LLC
//! ORCID: 0009-0005-2158-7060
//!
//! # Updated
//! April 16, 2026 11:15 PM EST
//!
//! # Disclaimer
//! Use at own risk. This software is provided "as is", without warranty of any
//! kind, express or implied. The author and AutomataNexus shall not be held
//! liable for any damages arising from the use of this software.

use axonml::prelude::*;

// =============================================================================
// Main Entry Point
// =============================================================================

fn main() {
    println!("=== Axonml ML Framework - Simple Training Example ===\n");

    // Print version and features
    println!("Version: {}", axonml::version());
    println!("Features: {}\n", axonml::features());

    // -------------------------------------------------------------------------
    // Dataset (XOR Problem)
    // -------------------------------------------------------------------------

    // 1. Create a simple dataset (XOR problem)
    println!("1. Creating XOR dataset...");
    let inputs = vec![
        vec![0.0, 0.0],
        vec![0.0, 1.0],
        vec![1.0, 0.0],
        vec![1.0, 1.0],
    ];
    let targets = vec![0.0, 1.0, 1.0, 0.0]; // XOR outputs

    println!("   Inputs: {inputs:?}");
    println!("   Targets: {targets:?}\n");

    // -------------------------------------------------------------------------
    // Model and Optimizer
    // -------------------------------------------------------------------------

    // 2. Create a simple MLP model
    println!("2. Creating MLP model (2 -> 4 -> 1)...");
    let linear1 = Linear::new(2, 4);
    let linear2 = Linear::new(4, 1);

    println!("   Layer 1: Linear(2, 4)");
    println!("   Layer 2: Linear(4, 1)\n");

    // 3. Create optimizer
    println!("3. Creating Adam optimizer (lr=0.1)...");
    let params = [linear1.parameters(), linear2.parameters()].concat();
    let mut optimizer = Adam::new(params, 0.1);
    println!("   Optimizer created!\n");

    // -------------------------------------------------------------------------
    // Training Loop
    // -------------------------------------------------------------------------

    // 4. Training loop
    println!("4. Training for 1000 epochs...");
    let epochs = 1000;

    for epoch in 0..epochs {
        let mut total_loss = 0.0;

        for (input, &target) in inputs.iter().zip(targets.iter()) {
            // Create input tensor
            let x = Variable::new(Tensor::from_vec(input.clone(), &[1, 2]).unwrap(), true);

            // Forward pass
            let h = linear1.forward(&x);
            let h = h.sigmoid();
            let output = linear2.forward(&h);
            let output = output.sigmoid();

            // Create target tensor
            let y = Variable::new(Tensor::from_vec(vec![target], &[1, 1]).unwrap(), false);

            // Compute MSE loss manually: (output - target)^2
            let diff = output.sub_var(&y);
            let loss = diff.mul_var(&diff);

            total_loss += loss.data().to_vec()[0];

            // Backward pass
            loss.backward();

            // Update weights
            optimizer.step();
            optimizer.zero_grad();
        }

        if epoch % 200 == 0 || epoch == epochs - 1 {
            println!("   Epoch {}: Loss = {:.6}", epoch, total_loss / 4.0);
        }
    }

    // -------------------------------------------------------------------------
    // Evaluation
    // -------------------------------------------------------------------------

    // 5. Test the trained model
    println!("\n5. Testing trained model...");
    for (input, &expected) in inputs.iter().zip(targets.iter()) {
        let x = Variable::new(Tensor::from_vec(input.clone(), &[1, 2]).unwrap(), false);

        let h = linear1.forward(&x);
        let h = h.sigmoid();
        let output = linear2.forward(&h);
        let output = output.sigmoid();

        let pred = output.data().to_vec()[0];
        let rounded = if pred > 0.5 { 1.0 } else { 0.0 };

        println!(
            "   Input: {input:?} -> Predicted: {pred:.4} (rounded: {rounded}) | Expected: {expected}"
        );
    }

    println!("\n=== Training Complete! ===");
}