Learning XOR

// Aner Neural: train a small network on the four synthetic XOR cases.
import aner.nn;

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);
        // 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);
    print("Final MSE");
    print(nn.item(nn.mse(prediction, y)));
    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));
}