test_mnist_01.ml

open Owl
open 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 l0 = {
  w = Maths.(Mat.uniform 784 300 * F 0.15 - F 0.075);
  b = Mat.zeros 1 300;
  a = Maths.tanh;
}

let l1 = {
  w = Maths.(Mat.uniform 300 10 * F 0.15 - F 0.075);
  b = Mat.zeros 1 10;
  a = Mat.map_by_row Maths.softmax;
}

let nn = {layers = [|l0; l1|]}