Iris: the full training loop

// Iris classification: native Aner, fixed stratified split, training-only scaling.
// Data: Fisher, R. (1936). Iris. UCI. https://doi.org/10.24432/C56C76
// CC BY 4.0: https://creativecommons.org/licenses/by/4.0/
// Original UCI iris.data retained; runtime class names map to zero-based IDs.
// Full attribution and variant notes: datasets/iris/README.md.
import aner.dataset;
import aner.nn;
import aner.metrics;
import aner.viz;

fn logits(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.add_bias(nn.matmul(hidden, w2), b2);
}

fn main() -> Unit {
    let flowers = dataset.iris();
    let split_seed = 2026;
    let partitions = dataset.split(flowers, 0.2, split_seed);
    let training = dataset.train(partitions);
    let testing = dataset.test(partitions);

    // Fit on training rows only; reuse those exact statistics for the test set.
    let scaler = dataset.fit_standardizer(dataset.features(training));
    let x_train = dataset.transform(scaler, dataset.features(training));
    let x_test = dataset.transform(scaler, dataset.features(testing));
    let y_train = dataset.one_hot(training);
    let train_labels = dataset.targets(training);
    let test_labels = dataset.targets(testing);

    var w1 = nn.parameter(nn.random(4, 8, 42, 0.5));
    var b1 = nn.parameter(nn.zeros(1, 8));
    var w2 = nn.parameter(nn.random(8, 3, 43, 0.5));
    var b2 = nn.parameter(nn.zeros(1, 3));
    let epochs = 2000;
    let learning_rate = 0.1;

    print("Dataset citation"); print(dataset.citation(flowers));
    print("Dataset license"); print(dataset.license(flowers));
    print("Dataset version"); print(dataset.version(flowers));
    print("Dataset SHA256"); print(dataset.sha256(flowers));
    print("Split seed"); print(split_seed);
    print("Training rows"); print(dataset.rows(training));
    print("Test rows"); print(dataset.rows(testing));
    print("Epochs"); print(epochs);
    print("Learning rate"); print(learning_rate);
    print("Initial training cross-entropy");
    print(nn.item(nn.cross_entropy(logits(x_train, w1, b1, w2, b2), y_train)));

    var epoch = 0;
    while epoch < epochs {
        let scores = logits(x_train, w1, b1, w2, b2);
        // Pass raw logits to this fused, numerically stable loss.
        let loss = nn.cross_entropy(scores, y_train);
        nn.backward(loss);
        if epoch < 3 || epoch % 100 == 0 { viz.capture(epoch, loss); }
        w1 = nn.sgd(w1, learning_rate);
        b1 = nn.sgd(b1, learning_rate);
        w2 = nn.sgd(w2, learning_rate);
        b2 = nn.sgd(b2, learning_rate);
        epoch = epoch + 1;
    }

    let train_scores = logits(x_train, w1, b1, w2, b2);
    let final_loss = nn.cross_entropy(train_scores, y_train);
    nn.backward(final_loss);
    viz.capture(epochs, final_loss);
    print("Final training cross-entropy"); print(nn.item(final_loss));
    print("Training accuracy"); print(metrics.accuracy(nn.argmax(train_scores), train_labels));

    // Evaluate the held-out rows once after the fixed training schedule.
    let test_scores = logits(x_test, w1, b1, w2, b2);
    print("Test accuracy"); print(metrics.accuracy(nn.argmax(test_scores), test_labels));
    let probabilities = nn.softmax(test_scores);
    let predictions = nn.argmax(probabilities);
    print("First test row class scores (probabilities)");
    var class_id = 0;
    while class_id < dataset.classes(flowers) {
        print(dataset.class_name(flowers, class_id));
        print(nn.value(probabilities, 0, class_id));
        class_id = class_id + 1;
    }
    print("First test row predicted class ID"); print(nn.value(predictions, 0, 0));
    print("First test row true class ID"); print(nn.value(test_labels, 0, 0));
}