using owl to build a trivial nn sing AD module from scratch

owl_nn.ml

open Owl_algodiff_ad

type layer = { mutable w : t; mutable b : t; a : t -> t }

type network = { layers : layer array }

let run_layer x l = Maths.((x $@ l.w) + l.b) |> l.a

let run_network x nn = Array.fold_left run_layer x nn.layers

let backprop nn eta epoch x y =
  let t = tag () in
  for i = 1 to epoch do
    Array.iter (fun l ->
      l.w <- make_reverse l.w t;
      l.b <- make_reverse l.b t;
    ) nn.layers;
    Maths.(l2norm_sqr((run_network x nn) - y)) |> reverse_prop (F 1.);
    Array.iter (fun l ->
      l.w <- Maths.((primal l.w) - (eta * (adjval l.w))) |> primal;
      l.b <- Maths.((primal l.b) - (eta * (adjval l.b))) |> primal;
    ) nn.layers;
  done