See a network learn XOR

// Aner Visual Debugger: record the same synthetic XOR training without changing updates.
// Run with --debug-viz NEW_DIRECTORY; add --debug-values for the synthetic neuron view.
import aner.nn;
import aner.viz;

fn predict(x: Tensor, w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> Tensor {
    let hidden = nn.tanh(nn.add_bias(nn.matmul(x, w1), b1));
    return nn.sigmoid(nn.add_bias(nn.matmul(hidden, w2), b2));
}

fn main() -> Unit {
    let x: Tensor = [[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]];
    let y: Tensor = [[0.0], [1.0], [1.0], [0.0]];

    var w1: Tensor = nn.parameter(nn.random(2, 8, 42, 1.0));
    var b1: Tensor = nn.parameter(nn.zeros(1, 8));
    var w2: Tensor = nn.parameter(nn.random(8, 1, 43, 1.0));
    var b2: Tensor = nn.parameter(nn.zeros(1, 1));

    print("Initial MSE");
    print(nn.item(nn.mse(predict(x, w1, b1, w2, b2), y)));

    var step = 0;
    while step < 3000 {
        let prediction = predict(x, w1, b1, w2, b2);
        let loss = nn.mse(prediction, y);
        nn.backward(loss);
        // Step counts completed updates. Capture before the next update.
        if step < 3 || step % 100 == 0 {
            viz.capture(step, loss);
        }
        // Each update returns a new leaf; assign it for the next iteration.
        w1 = nn.sgd(w1, 0.5);
        b1 = nn.sgd(b1, 0.5);
        w2 = nn.sgd(w2, 0.5);
        b2 = nn.sgd(b2, 0.5);
        step = step + 1;
    }

    let prediction = predict(x, w1, b1, w2, b2);
    let final_loss = nn.mse(prediction, y);
    nn.backward(final_loss);
    viz.capture(step, final_loss);
    print("Final MSE");
    print(nn.item(final_loss));
    print("Predictions for [0,0], [0,1], [1,0], [1,1]");
    print(nn.value(prediction, 0, 0));
    print(nn.value(prediction, 1, 0));
    print(nn.value(prediction, 2, 0));
    print(nn.value(prediction, 3, 0));
}