feat: MNIST

build.zig

@@ -94,7 +94,7 @@ pub fn build(b: *std.Build) void {
     // This will evaluate the `run` step rather than the default step.
     // For a top level step to actually do something, it must depend on other
     // steps (e.g. a Run step, as we will see in a moment).
-    const run_step = b.step("run", "Run the app");
+    const run_step = b.step("demo", "Run the engine autograd demo");
 
     // This creates a RunArtifact step in the build graph. A RunArtifact step
     // invokes an executable compiled by Zig. Steps will only be executed by the
@@ -115,6 +115,27 @@ pub fn build(b: *std.Build) void {
         run_cmd.addArgs(args);
     }
 
+    // The MNIST training example. Forced to ReleaseFast because scalar autograd
+    // over thousands of examples is painfully slow in Debug. `zig build run`
+    // builds and runs it; `mnist.zig` reads the CSVs from `mnist/` (run from the
+    // project root, which is the cwd `zig build` uses).
+    const mnist_exe = b.addExecutable(.{
+        .name = "mnist",
+        .root_module = b.createModule(.{
+            .root_source_file = b.path("src/mnist.zig"),
+            .target = target,
+            .optimize = .ReleaseFast,
+        }),
+    });
+    b.installArtifact(mnist_exe);
+
+    const run_mnist = b.addRunArtifact(mnist_exe);
+    run_mnist.step.dependOn(b.getInstallStep());
+    if (b.args) |args| run_mnist.addArgs(args);
+
+    const mnist_step = b.step("run", "Train & test an MLP on MNIST");
+    mnist_step.dependOn(&run_mnist.step);
+
     // Creates an executable that will run `test` blocks from the provided module.
     // Here `mod` needs to define a target, which is why earlier we made sure to
     // set the releative field.

src/mnist.zig

@@ -0,0 +1,187 @@
+const std = @import("std");
+const print = std.debug.print;
+
+const engine = @import("engine.zig");
+const nn = @import("nn.zig");
+
+// MNIST classifier trained with the scalar autograd engine. Because every op
+// allocates a Value into one static arena that only frees on deinit, a 784->H->10
+// forward pass mints ~50k Values per example — far too much to keep around for a
+// whole epoch. So weights live as plain []f32 and the Value graph is rebuilt and
+// thrown away every training step (see the train loop). This keeps memory flat.
+
+// Tunables — kept small so `zig build run` finishes in seconds, not hours.
+const TRAIN_N = 2000;
+const TEST_N = 1000;
+const HIDDEN = 32;
+const EPOCHS = 10;
+const LR: f32 = 0.05;
+
+const PIXELS = 784; // 28 * 28
+const CLASSES = 10;
+
+const SIZES = [_]usize{ PIXELS, HIDDEN, CLASSES };
+const Net = nn.MLP(f32, &SIZES);
+const Value = engine.Value(f32);
+
+const Data = struct {
+    n: usize,
+    labels: []u8, // n
+    pixels: []f32, // n * PIXELS, normalised to [0, 1]
+};
+
+/// Read the first `want` rows of a label-first MNIST CSV (no header).
+fn loadCsv(io: std.Io, gpa: std.mem.Allocator, path: []const u8, want: usize) !Data {
+    const raw = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited);
+    defer gpa.free(raw);
+
+    const labels = try gpa.alloc(u8, want);
+    const pixels = try gpa.alloc(f32, want * PIXELS);
+
+    var lines = std.mem.tokenizeScalar(u8, raw, '\n');
+    var row: usize = 0;
+    while (row < want) {
+        const line = lines.next() orelse break;
+        if (line.len == 0) continue;
+
+        var cols = std.mem.tokenizeScalar(u8, line, ',');
+        const lab = cols.next() orelse continue;
+        labels[row] = try std.fmt.parseInt(u8, std.mem.trim(u8, lab, " \r"), 10);
+
+        var c: usize = 0;
+        while (cols.next()) |tok| : (c += 1) {
+            if (c >= PIXELS) break;
+            const v = try std.fmt.parseInt(u16, std.mem.trim(u8, tok, " \r"), 10);
+            pixels[row * PIXELS + c] = @as(f32, @floatFromInt(v)) / 255.0;
+        }
+        row += 1;
+    }
+    return .{ .n = row, .labels = labels, .pixels = pixels };
+}
+
+/// Xavier-ish init: weights ~ U(-1, 1) / sqrt(fan_in), biases 0. Without the
+/// 1/sqrt(fan_in) scale, 784 inputs blow preactivations up and the loss diverges.
+/// Layout matches MLP.parameters()/fromParams: per layer, per neuron, weights then bias.
+fn initParams(params: []f32, seed: u64) void {
+    var rng = nn.Random(f32).init(seed);
+    var off: usize = 0;
+    inline for (0..SIZES.len - 1) |l| {
+        const fan_in = SIZES[l];
+        const fan_out = SIZES[l + 1];
+        const scale = 1.0 / @sqrt(@as(f32, @floatFromInt(fan_in)));
+        for (0..fan_out) |_| {
+            for (0..fan_in) |_| {
+                params[off] = (rng.uniform() * 2 - 1) * scale;
+                off += 1;
+            }
+            params[off] = 0; // bias
+            off += 1;
+        }
+    }
+}
+
+/// Fraction of `data` classified correctly (argmax of the logits).
+fn accuracy(gpa: std.mem.Allocator, params: []const f32, data: Data) f32 {
+    var correct: usize = 0;
+    var i: usize = 0;
+    while (i < data.n) : (i += 1) {
+        var net = Net.fromParams(gpa, params);
+        defer net.deinit();
+
+        var x: [PIXELS]*Value = undefined;
+        for (0..PIXELS) |j| x[j] = Value.new(data.pixels[i * PIXELS + j]);
+        const logits = net.forward(x);
+
+        var best: usize = 0;
+        var best_v: f32 = logits[0].data;
+        for (1..CLASSES) |k| {
+            if (logits[k].data > best_v) {
+                best_v = logits[k].data;
+                best = k;
+            }
+        }
+        if (best == data.labels[i]) correct += 1;
+    }
+    return @as(f32, @floatFromInt(correct)) / @as(f32, @floatFromInt(data.n));
+}
+
+pub fn main() !void {
+    var gpa_state = std.heap.DebugAllocator(.{}){};
+    defer _ = gpa_state.deinit();
+    const gpa = gpa_state.allocator();
+
+    // 0.16 file IO goes through an Io instance.
+    var threaded = std.Io.Threaded.init(gpa, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
+
+    print("loading MNIST...\n", .{});
+    const train = try loadCsv(io, gpa, "mnist/mnist_train.csv", TRAIN_N);
+    defer {
+        gpa.free(train.labels);
+        gpa.free(train.pixels);
+    }
+    const test_set = try loadCsv(io, gpa, "mnist/mnist_test.csv", TEST_N);
+    defer {
+        gpa.free(test_set.labels);
+        gpa.free(test_set.pixels);
+    }
+    print("train={d} test={d}  net={d}->{d}->{d}  params={d}\n", .{
+        train.n, test_set.n, PIXELS, HIDDEN, CLASSES, Net.param_count,
+    });
+
+    // Persistent weights (survive the per-step arena resets).
+    const params = try gpa.alloc(f32, Net.param_count);
+    defer gpa.free(params);
+    initParams(params, 1);
+
+    // Reusable topological-sort backprop (its buffers live on gpa, not the arena).
+    var bp = engine.Backprop(Value).init(gpa);
+    defer bp.deinit();
+
+    // Shuffle order each epoch.
+    var idx = try gpa.alloc(usize, train.n);
+    defer gpa.free(idx);
+    for (0..train.n) |i| idx[i] = i;
+    var prng = std.Random.DefaultPrng.init(0xC0FFEE);
+    const rand = prng.random();
+
+    print("\ntraining ({d} epochs, lr={d})\n", .{ EPOCHS, LR });
+    for (0..EPOCHS) |epoch| {
+        rand.shuffle(usize, idx);
+        var total_loss: f32 = 0;
+
+        for (idx) |i| {
+            // Fresh arena + weights for this example.
+            var net = Net.fromParams(gpa, params);
+
+            var x: [PIXELS]*Value = undefined;
+            for (0..PIXELS) |j| x[j] = Value.new(train.pixels[i * PIXELS + j]);
+            const logits = net.forward(x);
+
+            // Sum-of-squares loss against the one-hot target.
+            var loss = Value.new(0);
+            const label = train.labels[i];
+            for (0..CLASSES) |k| {
+                const t: f32 = if (k == label) 1.0 else 0.0;
+                loss = loss.add(logits[k].sub(t).pow(2));
+            }
+            total_loss += loss.data;
+
+            bp.backprop(loss);
+
+            // SGD: copy grads out, then drop the whole graph.
+            const ps = net.parameters();
+            for (0..Net.param_count) |k| params[k] -= LR * ps[k].grad;
+
+            net.deinit();
+        }
+
+        const acc = accuracy(gpa, params, test_set);
+        print("epoch {d}: avg loss {d:.4}  test acc {d:.3}\n", .{
+            epoch, total_loss / @as(f32, @floatFromInt(train.n)), acc,
+        });
+    }
+
+    print("\nfinal test accuracy: {d:.3}\n", .{accuracy(gpa, params, test_set)});
+}

src/nn.zig

@@ -47,6 +47,15 @@ pub fn Neuron(comptime T: type, comptime in: usize) type {
             for (&w) |*wi| wi.* = Value.new(rng.uniform() * 2 - 1);
             return .{ .weights = w, .bias = Value.new(0) };
         }
+
+        /// Build a neuron from saved parameters (weights then bias), consuming
+        /// `param_count` entries from the front of `buf`. Pairs with `parameters`.
+        pub fn fromParams(buf: []const T) Self {
+            var w: [in]*Value = undefined;
+            inline for (0..in) |i| w[i] = Value.new(buf[i]);
+            return .{ .weights = w, .bias = Value.new(buf[in]) };
+        }
+
         /// w·x + b, optionally passed through ReLU. `nonlin` is comptime so the
         /// branch is resolved away — output neurons compile to the linear path.
         pub fn forward(self: *const Self, x: [in]*Value, comptime nonlin: bool) *Value {
@@ -83,6 +92,13 @@ pub fn Layer(comptime T: type, comptime in: usize, comptime out: usize) type {
             return .{ .neurons = ns };
         }
 
+        /// Build a layer from saved parameters, consuming `param_count` entries.
+        pub fn fromParams(buf: []const T) Self {
+            var ns: [out]N = undefined;
+            inline for (0..out) |i| ns[i] = N.fromParams(buf[i * N.param_count ..]);
+            return .{ .neurons = ns };
+        }
+
         pub fn forward(self: *const Self, x: [in]*Value, comptime nonlin: bool) [out]*Value {
             var y: [out]*Value = undefined;
             inline for (&self.neurons, 0..) |*n, i| y[i] = n.forward(x, nonlin);
@@ -144,6 +160,21 @@ pub fn MLP(comptime T: type, comptime sizes: []const usize) type {
             return self;
         }
 
+        /// Build the net in a fresh Value arena from a saved parameter array
+        /// (same layout as `parameters()`). Used by training loops that reset the
+        /// arena each step: keep weights as plain `[]T`, rematerialise per step.
+        pub fn fromParams(alloc: std.mem.Allocator, data: []const T) Self {
+            Value.init(alloc);
+            var self: Self = undefined;
+            var off: usize = 0;
+            inline for (0..n_layers) |i| {
+                const LT = @TypeOf(self.layers[i]);
+                self.layers[i] = LT.fromParams(data[off..]);
+                off += LT.param_count;
+            }
+            return self;
+        }
+
         /// Tears down the engine Value arena (frees every Value ever created).
         pub fn deinit(self: *Self) void {
             _ = self;