From dca0efd1f6f0da12825f955ef1d12affdb06f73a Mon Sep 17 00:00:00 2001
From: Mattia Giambirtone <nocturn9x@nocturn9x.space>
Date: Mon, 20 Mar 2023 10:31:09 +0100
Subject: [PATCH] Derivatives are now vectorized

---
 src/main.nim       | 4 ++--
 src/nn/network.nim | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/main.nim b/src/main.nim
index 3333cb5..007bb0b 100644
--- a/src/main.nim
+++ b/src/main.nim
@@ -10,9 +10,9 @@ proc mse(a, b: Matrix[float]): float =
     result = (b - a).apply(proc (x: float): float = pow(x, 2), axis = -1).sum() / len(a).float
 
 # Derivative of MSE
-func dxMSE*(x, y: float): float = 2 * (x - y)
+func dxMSE*(x, y: Matrix[float]): Matrix[float] = 2.0 * (x - y)
 
-func dx*(x, y: float): float = 0.0
+func dx*(x, y: Matrix[float]): Matrix[float] = zeros[float](x.shape)
 
 # A bunch of vectorized activation functions
 func sigmoid*(input: Matrix[float]): Matrix[float] = 
diff --git a/src/nn/network.nim b/src/nn/network.nim
index 5623838..e66b1b5 100644
--- a/src/nn/network.nim
+++ b/src/nn/network.nim
@@ -41,11 +41,11 @@ type
     Loss* = ref object
         ## A loss function and its derivative
         function: proc (a, b: Matrix[float]): float
-        derivative: proc (x, y: float): float {.noSideEffect.}
+        derivative: proc (x, y: Matrix[float]): Matrix[float] {.noSideEffect.}
     Activation* = ref object
         ## An activation function
         function: proc (input: Matrix[float]): Matrix[float] {.noSideEffect.}
-        derivative: proc (x, y: float): float {.noSideEffect.}
+        derivative: proc (x, y: Matrix[float]): Matrix[float] {.noSideEffect.}
     Layer* = ref object
         ## A generic neural network
         ## layer
@@ -70,14 +70,14 @@ proc `$`*(self: NeuralNetwork): string =
     result = &"NeuralNetwork(learnRate={self.learnRate}, layers={self.layers})"
 
 
-proc newLoss*(function: proc (a, b: Matrix[float]): float, derivative: proc (x, y: float): float {.noSideEffect.}): Loss =
+proc newLoss*(function: proc (a, b: Matrix[float]): float, derivative: proc (x, y: Matrix[float]): Matrix[float] {.noSideEffect.}): Loss =
     ## Creates a new Loss object
     new(result)
     result.function = function
     result.derivative = derivative
 
 
-proc newActivation*(function: proc (input: Matrix[float]): Matrix[float] {.noSideEffect.}, derivative: proc (x, y: float): float {.noSideEffect.}): Activation =
+proc newActivation*(function: proc (input: Matrix[float]): Matrix[float] {.noSideEffect.}, derivative: proc (x, y: Matrix[float]): Matrix[float] {.noSideEffect.}): Activation =
     ## Creates a new Activation object
     new(result)
     result.function = function