feat: implement nn

README.md

@@ -2,7 +2,7 @@
 
 Implementing Karpathy's [micrograd](https://github.com/karpathy/micrograd) with some help from [0xdeb7ef/microgradz](https://github.com/0xdeb7ef/microgradz).
 
-Scalar only implementation
+Scalar only implementation, the `nn.zig` module implements `Neuron`, `Layer` and `MLP` to have comptime-known shapes, only the `engine.zig` module needs runtime-allocated memory.
 
 ## Setup
 

src/main.zig

@@ -1,47 +1,66 @@
 const std = @import("std");
 const print = std.debug.print;
 
-// import Value and Backprop
-const ValueType = @import("engine.zig").Value;
-const BackpropType = @import("engine.zig").Backprop;
+const engine = @import("engine.zig");
+const nn = @import("nn.zig");
+
+const ValueType = engine.Value;
+const BackpropType = engine.Backprop;
 
 pub fn main() !void {
     var gpa = std.heap.DebugAllocator(.{}){};
     const alloc = gpa.allocator();
 
-    // Create the Value pool with a specific float type
     const Value = ValueType(f32);
-    Value.init(alloc); // init with an allocator, arena is used internally
-    defer Value.deinit();
 
-    // Create a reusable Backprop instance that operates on Value
-    var Backprop = BackpropType(Value).init(alloc); // init with an allocator
-    defer Backprop.deinit();
+    // The MLP owns the Value arena: init() calls Value.init, deinit() frees it.
+    const Net = nn.MLP(f32, &.{ 3, 4, 4, 1 });
+    var net = Net.init(alloc);
+    defer net.deinit();
 
-    var a = Value.new(-4);
-    var b = Value.new(2);
-    var c = a.add(b);
-    var d = a.mul(b).add(b.pow(3));
-    // c += c + 1
-    c = c.add(c.add(1));
-    // c += 1 + c + (-a)
-    c = c.add(c.add(1).sub(a));
-    // d += d * 2 + (b + a).relu()
-    d = d.add(d.mul(2).add(b.add(a).relu()));
-    // d += 3 * d + (b - a).relu()
-    d = d.add(d.mul(3).add(b.sub(a).relu()));
+    var backprop = BackpropType(Value).init(alloc);
+    defer backprop.deinit();
 
-    var e = c.sub(d);
-    var f = e.pow(2);
-    var g = f.div(2);
-    // g += 10.0 / f
-    g = g.add(Value.new(10).div(f));
+    // Karpathy's tiny binary dataset (targets are -1 / +1).
+    const xs = [_][3]f32{
+        .{ 2, 3, -1 },
+        .{ 3, -1, 0.5 },
+        .{ 0.5, 1, 1 },
+        .{ 1, 1, -1 },
+    };
+    const ys = [_]f32{ 1, -1, -1, 1 };
 
-    print("g.data: {d:.4}\n", .{g.data}); // 24.7041
+    const lr: f32 = 0.05;
+    var epoch: usize = 0;
+    while (epoch < 100) : (epoch += 1) {
+        // Build a single summed squared-error loss graph over the batch.
+        var loss = Value.new(0);
+        for (xs, ys) |x, y| {
+            var xv: [3]*Value = undefined;
+            for (x, 0..) |xi, i| xv[i] = Value.new(xi);
+            const pred = net.forward(xv);
+            loss = loss.add(pred[0].sub(y).pow(2));
+        }
 
-    // run backprop on `g`
-    Backprop.backprop(g);
+        net.zeroGrad();
+        backprop.backprop(loss);
+        for (net.parameters()) |p| p.data -= lr * p.grad;
 
-    print("a.grad: {d:.4}\n", .{a.grad}); // 138.8338
-    print("b.grad: {d:.4}\n", .{b.grad}); // 645.5773
+        if (epoch % 10 == 0) {
+            print("epoch {d:>3}: loss {d:.6}\n", .{ epoch, loss.data });
+        }
+    }
+
+    // Final predictions should track the targets (+/-1).
+    print("\npredictions:\n", .{});
+    for (xs, ys) |x, y| {
+        var xv: [3]*Value = undefined;
+        for (x, 0..) |xi, i| xv[i] = Value.new(xi);
+        print("  pred {d:>7.4}  target {d:>5.1}\n", .{ net.forward(xv)[0].data, y });
+    }
+}
+
+test {
+    _ = @import("engine.zig");
+    _ = @import("nn.zig");
 }

src/nn.zig

@@ -0,0 +1,221 @@
+const std = @import("std");
+const engine = @import("engine.zig");
+
+// Karpathy's micrograd `nn` module, ported to Zig with the network shape known
+// at comptime. The layer sizes are a comptime `[]const usize`, so every
+// structural array (weights, neurons, layers, parameter list) has a fixed,
+// comptime-known length. The ONLY runtime allocation happens inside engine's
+// `Value`, which keeps the whole autograd graph in a static arena.
+
+/// Small PRNG wrapper that yields uniform values in the unit interval (floats)
+/// or the full range (ints). Held by value so each `MLP` seeds its own stream.
+pub fn Random(comptime T: type) type {
+    return struct {
+        const Self = @This();
+
+        prng: std.Random.DefaultPrng,
+
+        pub fn init(seed: u64) Self {
+            return .{ .prng = std.Random.DefaultPrng.init(seed) };
+        }
+
+        pub fn uniform(self: *Self) T {
+            return switch (@typeInfo(T)) {
+                .float => self.prng.random().float(T),
+                .int => self.prng.random().int(T),
+                else => @compileError("Only integer and floating-point types are supported."),
+            };
+        }
+    };
+}
+
+/// A single neuron with `in` weights and a bias. Weights are initialised
+/// uniformly in (-1, 1) (Karpathy's default); the bias starts at 0.
+pub fn Neuron(comptime T: type, comptime in: usize) type {
+    const Value = engine.Value(T);
+    return struct {
+        const Self = @This();
+
+        /// Number of `*Value` parameters this neuron contributes (weights + bias).
+        pub const param_count = in + 1;
+
+        weights: [in]*Value,
+        bias: *Value,
+
+        pub fn init(rng: *Random(T)) Self {
+            var w: [in]*Value = undefined;
+            for (&w) |*wi| wi.* = Value.new(rng.uniform() * 2 - 1);
+            return .{ .weights = w, .bias = Value.new(0) };
+        }
+        /// 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 {
+            var sum = self.bias;
+            inline for (self.weights, x) |w, xi| {
+                sum = sum.add(w.mul(xi));
+            }
+            return if (nonlin) sum.relu() else sum;
+        }
+
+        /// Writes this neuron's parameters into `buf` and returns the count used.
+        pub fn parameters(self: *const Self, buf: []*Value) usize {
+            inline for (self.weights, 0..) |w, i| buf[i] = w;
+            buf[in] = self.bias;
+            return param_count;
+        }
+    };
+}
+
+/// A dense layer: `out` neurons, each taking `in` inputs.
+pub fn Layer(comptime T: type, comptime in: usize, comptime out: usize) type {
+    const Value = engine.Value(T);
+    const N = Neuron(T, in);
+    return struct {
+        const Self = @This();
+
+        pub const param_count = out * N.param_count;
+
+        neurons: [out]N,
+
+        pub fn init(rng: *Random(T)) Self {
+            var ns: [out]N = undefined;
+            for (&ns) |*n| n.* = N.init(rng);
+            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);
+            return y;
+        }
+
+        pub fn parameters(self: *const Self, buf: []*Value) usize {
+            var n: usize = 0;
+            inline for (&self.neurons) |*neu| n += neu.parameters(buf[n..]);
+            return n;
+        }
+    };
+}
+
+/// A multi-layer perceptron. `sizes[0]` is the input width; each subsequent
+/// entry is a layer's output width. Hidden layers use ReLU; the final layer is
+/// linear (Karpathy's `nonlin = i != len-1`).
+///
+///     const Net = MLP(f32, &.{ 3, 4, 4, 1 });
+///     var net = Net.init(alloc);   // owns the engine Value arena
+///     defer net.deinit();
+///     const out = net.forward(.{ x0, x1, x2 }); // [1]*Value
+pub fn MLP(comptime T: type, comptime sizes: []const usize) type {
+    if (sizes.len < 2) @compileError("MLP needs at least an input size and one layer.");
+
+    const Value = engine.Value(T);
+    const n_layers = sizes.len - 1;
+
+    // Layers have heterogeneous sizes, so they can't share a uniform array
+    // type. A comptime-built tuple holds one concrete `Layer(...)` per slot.
+    const LayersTuple = blk: {
+        var lt: [n_layers]type = undefined;
+        for (0..n_layers) |i| lt[i] = Layer(T, sizes[i], sizes[i + 1]);
+        break :blk std.meta.Tuple(&lt);
+    };
+
+    return struct {
+        const Self = @This();
+
+        pub const in_size = sizes[0];
+        pub const out_size = sizes[sizes.len - 1];
+        pub const param_count = blk: {
+            var c: usize = 0;
+            for (0..n_layers) |i| c += sizes[i + 1] * (sizes[i] + 1);
+            break :blk c;
+        };
+
+        layers: LayersTuple,
+
+        /// Initialise the engine Value arena and build weights.
+        pub fn init(alloc: std.mem.Allocator) Self {
+            const seed = 42;
+            Value.init(alloc);
+            var rng = Random(T).init(seed);
+            var self: Self = undefined;
+            inline for (0..n_layers) |i| {
+                self.layers[i] = @TypeOf(self.layers[i]).init(&rng);
+            }
+            return self;
+        }
+
+        /// Tears down the engine Value arena (frees every Value ever created).
+        pub fn deinit(self: *Self) void {
+            _ = self;
+            Value.deinit();
+        }
+
+        pub fn forward(self: *const Self, x: [in_size]*Value) [out_size]*Value {
+            return self.forwardFrom(0, x);
+        }
+
+        // Comptime recursion threads the typed activation array through layers;
+        // `i` being comptime is what lets the array type shrink/grow per layer.
+        fn forwardFrom(self: *const Self, comptime i: usize, x: [sizes[i]]*Value) [out_size]*Value {
+            const nonlin = i != n_layers - 1;
+            const y = self.layers[i].forward(x, nonlin);
+            if (i == n_layers - 1) return y;
+            return self.forwardFrom(i + 1, y);
+        }
+
+        /// All trainable parameters as a fixed-size array (no allocation).
+        pub fn parameters(self: *const Self) [param_count]*Value {
+            var buf: [param_count]*Value = undefined;
+            var n: usize = 0;
+            inline for (0..n_layers) |i| n += self.layers[i].parameters(buf[n..]);
+            return buf;
+        }
+
+        pub fn zeroGrad(self: *const Self) void {
+            for (self.parameters()) |p| p.grad = 0;
+        }
+    };
+}
+
+test "param_count and parameters() length agree" {
+    const Net = MLP(f32, &.{ 3, 4, 4, 1 });
+    var net = Net.initSeed(std.testing.allocator, 42);
+    defer net.deinit();
+
+    // 4*(3+1) + 4*(4+1) + 1*(4+1) = 16 + 20 + 5 = 41
+    try std.testing.expectEqual(@as(usize, 41), Net.param_count);
+    try std.testing.expectEqual(@as(usize, 41), net.parameters().len);
+}
+
+test "forward produces the output-layer width" {
+    const Value = engine.Value(f32);
+    const Net = MLP(f32, &.{ 3, 4, 4, 1 });
+    var net = Net.initSeed(std.testing.allocator, 7);
+    defer net.deinit();
+
+    const x: [3]*Value = .{ Value.new(1), Value.new(-2), Value.new(3) };
+    const y = net.forward(x);
+    try std.testing.expectEqual(@as(usize, 1), y.len);
+}
+
+test "one gradient-descent step lowers the loss" {
+    const Value = engine.Value(f32);
+    const Backprop = engine.Backprop(Value);
+    const Net = MLP(f32, &.{ 2, 8, 1 });
+    var net = Net.initSeed(std.testing.allocator, 123);
+    defer net.deinit();
+    var bp = Backprop.init(std.testing.allocator);
+    defer bp.deinit();
+
+    const x: [2]*Value = .{ Value.new(0.5), Value.new(-0.5) };
+    const target: f32 = 1.0;
+
+    const loss0 = net.forward(x)[0].sub(target).pow(2);
+    net.zeroGrad();
+    bp.backprop(loss0);
+    const before = loss0.data;
+    for (net.parameters()) |p| p.data -= 0.01 * p.grad;
+
+    const after = net.forward(x)[0].sub(target).pow(2).data;
+    try std.testing.expect(after < before);
+}