Cleaner training script

Author: Kaweees

examples/train.zig

@@ -5,11 +5,6 @@
 const std = @import("std");
 /// This imports the separate module containing `root.zig`. Take a look in `build.zig` for details.
 const kiwigrad = @import("kiwigrad");
-const zbench = @import("zbench");
-
-const ValueType = kiwigrad.engine.Value;
-const NeuronType = kiwigrad.nn.Neuron;
-const LayerType = kiwigrad.nn.Layer;
 
 /// Write the computational graph to a Graphviz file
 pub fn draw_graph(comptime T: type, graph: *kiwigrad.engine.Value(T), name: []const u8, writer: anytype) !void {
@@ -33,24 +28,39 @@ pub fn main() !void {
     const stdout = bw.writer();
     const alloc = std.heap.page_allocator;
 
-    // Initialize the value allocator
-    const Value = ValueType(f64);
-    Value.init(alloc);
-    defer Value.deinit();
+    // Initialize the required components
+    const ValueType = kiwigrad.engine.Value(f64);
+    const NeuronType = kiwigrad.nn.Neuron(f64);
+    const LayerType = kiwigrad.nn.Layer(f64);
+    // const MLPType = kiwigrad.nn.MLP;
 
-    // Initialize the neuron allocator
-    const Neuron = NeuronType(f64);
-    Neuron.init(alloc);
-    defer Neuron.deinit();
+    // Initialize allocators and components
+    ValueType.init(alloc);
+    NeuronType.init(alloc);
+    LayerType.init(alloc);
+    defer {
+        ValueType.deinit();
+        NeuronType.deinit();
+        LayerType.deinit();
+        // MLPType.deinit();
+    }
 
     // Initialize the neuron
-    const n = Neuron.new(2);
+    const neuron = NeuronType.new(3);
+
     // Create sample input data
-    var input_data = [_]*kiwigrad.engine.Value(f64){
-        kiwigrad.engine.Value(f64).new(1.0),
-        kiwigrad.engine.Value(f64).new(2.0),
+    var input_data = [_]*ValueType{
+        ValueType.new(1.0),
+        ValueType.new(2.0),
+        ValueType.new(3.0),
     };
-    const output = n.forward(input_data[0..]);
+
+    // Forward pass through the layer
+    const output = neuron.forward(input_data[0..]);
+
+    // outputs now contains 2 ValueType pointers (one for each neuron)
+    std.debug.print("Layer output: {d}\n", .{output.data});
+
     try draw_graph(f64, output, "n_f64", stdout);
     try bw.flush(); // Don't forget to flush!
 }