Add custom operators; add simple tests

lib/dune

@@ -1,3 +1,4 @@
 (library
-  (name nn)
-  (modules Neuron Network))
+  (name smolgrad)
+  (public_name smolgrad)
+  (modules neuron))

lib/neuron.ml

@@ -5,19 +5,22 @@ module Neuron = struct
     mutable backward : unit -> unit;
 
     operator : string;
-    dependents : t list;
+    dependencies : t list;
   }
 
-  let create dt op deps = { 
+  let data base =
+    base.data
+
+  let create ?(op="n/a") ?(deps=[]) dt = { 
     data = dt;
     grad = 0.;
     backward = (fun () -> ());
     operator = op;
-    dependents = deps;
+    dependencies = deps;
   }
 
   let add base partner =
-    let resultant = create (base.data +. partner.data) "+" [base; partner] in
+    let resultant = create ~op:"+" ~deps:[base; partner] (base.data +. partner.data) in
 
     resultant.backward <- (fun () ->
       base.grad <- base.grad +. resultant.grad;
@@ -26,7 +29,7 @@ module Neuron = struct
     resultant
 
   let mul base partner =
-    let resultant = create (base.data *. partner.data) "*" [base; partner] in
+    let resultant = create ~op:"*" ~deps:[base; partner] (base.data *. partner.data) in
 
     resultant.backward <- (fun () ->
       base.grad <- base.grad +. partner.data *. resultant.grad;
@@ -35,15 +38,20 @@ module Neuron = struct
     resultant
 
   let exp base exponent =
-    let resultant = create (base.data ** exponent) "**" [base] in
+    let resultant = create ~op:"**" ~deps:[base] (base.data ** exponent) in
 
     resultant.backward <- (fun () ->
       base.grad <- base.grad +. exponent *. (base.data ** (exponent -. 1.)) *. resultant.grad;
     );
     resultant
 
+(* these are all the operator overloadings we need to associate with each of the binary operators *)
+  let ( + ) = add
+  let ( * ) = mul
+  let ( ** ) = exp
+
   let relu base =
-    let resultant = create (max 0. base.data) "relu" [base] in
+    let resultant = create ~op:"relu" ~deps:[base] (max 0. base.data) in
 
     resultant.backward <- (fun () ->
       base.grad <- base.grad +. (if base.data > 0. then resultant.grad else 0.);
@@ -58,7 +66,7 @@ module Neuron = struct
         let visited, resultant =
           List.fold_left (fun (visited, resultant) dependent ->
             sort_topologically visited resultant dependent
-          ) (visited, resultant) candidate.dependents 
+          ) (visited, resultant) candidate.dependencies 
         in
         (visited, candidate :: resultant)
       else
@@ -72,3 +80,4 @@ module Neuron = struct
     (* and propagate the gradient changes *)
     List.iter (fun v -> v.backward ()) (List.rev resultant)
 end
+

lib/neuron.mli

@@ -1,22 +1,28 @@
 module Neuron : sig
     type t
 
+    (* getter for the data *)
+    val data : t -> float
+
     (* Constructor; constructs a unit neuron of a value and an operator. *)
-    val create : float -> string -> t list -> t
+    val create : ?op:string -> ?deps:t list -> float -> t
 
     (* Handles the gradient flows in addition operation. *)
     val add : t -> t -> t
+    val ( + ) : t -> t -> t
 
     (* Handles the gradient flows in multiplication operation. *)
     val mul : t -> t -> t
+    val ( * ) : t -> t -> t
 
     (* Handles the gradient flows in exponent / power operation. *)
     (* second argument is the exponent. *)
     val exp : t -> float -> t
+    val ( ** ) : t -> float -> t
 
     (* Handles the gradient flows in ReLU operation. *)
     val relu : t -> t
 
-    (* Handles backpropagation of the gradients for all the nodes connected to the specified base node. *)
+    (* Handles backpropagation of the gradients for all the nodes connected to this as the base node. *)
     val backpropagate : t -> unit
 end

test/dune

@@ -0,0 +1,3 @@
+(test
+ (name smolgrad_tests)
+ (libraries ounit2 smolgrad))

test/smolgrad_tests.ml

@@ -0,0 +1,24 @@
+open Smolgrad.Neuron
+
+let test_simple_operation () =
+  let a = Neuron.create 4.0 in
+  let b = Neuron.create 2.0 in
+
+  let abba = Neuron.(a + b) in
+  assert (Neuron.data abba = 6.0);
+;;
+
+(* here we open the Neuron module wide open locally, thereby allowing the clean custom `+` operator usage *)
+let test_custom_operator () =
+  let open Neuron in
+  let a = create 4.0 in
+  let b = create 2.0 in
+
+  let abba = a + b in
+  assert (Neuron.data abba = 6.0);
+;;
+
+let () =
+  test_simple_operation ();
+  test_custom_operator ();
+;;