main.cpp

#include<iostream>
#include "tensor.h"
#include "sgd.h"
using namespace std;

int main() {
    // XOR dataset
    Tensor X({4, 2}, {
        0, 0,
        0, 1,
        1, 0,
        1, 1
    });

    Tensor Y({4, 1}, {
        0,
        1,
        1,
        0
    });

    // weights — small random init
    Tensor W1({2, 4}, {
         0.1f, -0.2f,  0.3f,  0.1f,
        -0.1f,  0.2f, -0.3f,  0.2f
    });

    Tensor W2({4, 1}, {
         0.2f,
        -0.1f,
         0.3f,
        -0.2f
    });

    SGD optimizer({&W1, &W2}, 0.01f);

    // training loop
    for (int epoch = 0; epoch <= 1000; epoch++) {

        // forward pass
        auto h    = matmul(X, W1);       // (4,2) @ (2,4) = (4,4)
        auto h2   = relu(*h);            // (4,4)
        auto out  = matmul(*h2, W2);     // (4,4) @ (4,1) = (4,1)
        auto loss = mse(*out, Y);        // scalar (1,)

        // print every 100 epochs
        if (epoch % 100 == 0) {
            cout << "epoch " << epoch << "  loss: " << loss->data[0] << endl;
        }

        // backward + update
        Tensor::backward(loss.get());
        optimizer.step();
        optimizer.zero_grad();
    }

    // final predictions
    cout << "\nFinal predictions:" << endl;
    auto h   = matmul(X, W1);
    auto h2  = relu(*h);
    auto out = matmul(*h2, W2);

    cout << "Input [0,0] -> " << out->data[0] << "  (expected 0)" << endl;
    cout << "Input [0,1] -> " << out->data[1] << "  (expected 1)" << endl;
    cout << "Input [1,0] -> " << out->data[2] << "  (expected 1)" << endl;
    cout << "Input [1,1] -> " << out->data[3] << "  (expected 0)" << endl;

    return 0;
}