Source file owl_neural_graph.ml

# 1 "src/base/neural/owl_neural_graph.ml"
(*
 * OWL - OCaml Scientific Computing
 * Copyright (c) 2016-2022 Liang Wang <liang@ocaml.xyz>
 *)

(** Neural network: Graphical neural network *)

open Owl_types

(* Make functor starts *)

module Make (Neuron : Owl_neural_neuron_sig.Sig) = struct
  module Neuron = Neuron
  open Neuron
  open Neuron.Optimise.Algodiff

  (* graph network and node definition *)

  type node =
    { mutable name : string
    ; (* name of a node *)
      mutable prev : node array
    ; (* parents of a node *)
      mutable next : node array
    ; (* children of a node *)
      mutable neuron : neuron
    ; (* neuron contained in a node *)
      mutable output : t option
    ; (* output of a node *)
      mutable network : network
    ; (* network a node belongs to *)
      mutable train : bool (* specify if a node is only for training *)
    }

  and network =
    { mutable nnid : string
    ; (* name of the graph network *)
      mutable size : int
    ; (* size of the graph network *)
      mutable roots : node array
    ; (* roots of the graph network, i.e. inputs *)
      mutable outputs : node array
    ; (* outputs of the graph network *)
      mutable topo : node array (* nodes sorted in topological order *)
    }

  (* functions to manipulate the network *)

  let make_network ?nnid size roots topo =
    let nnid =
      match nnid with
      | Some s -> s
      | None   -> "Graphical network"
    in
    { nnid; size; roots; topo; outputs = [||] }


  let make_node ?name ?(train = false) prev next neuron output network =
    let name =
      match name with
      | Some s -> s
      | None   -> Printf.sprintf "%s_%i" (to_name neuron) network.size
    in
    { name; prev; next; neuron; output; network; train }


  let get_roots nn =
    match nn.roots with
    | [||] -> failwith "Owl_neural_graph:get_roots"
    | x    -> x


  let get_outputs nn = nn.outputs

  let get_node nn name =
    let x = Owl_utils.Array.filter (fun n -> n.name = name) nn.topo in
    if Array.length x = 0 then failwith "Owl_neural_graph:get_node" else x.(0)


  let get_network ?name n =
    let name =
      match name with
      | Some s -> s
      | None   -> Random.int 65535 |> string_of_int
    in
    n.network.nnid <- name;
    (* if run and run_inputs are merged, the next line is necessary. *)
    n.network.outputs <- [| n |];
    n.network


  let outputs ?name nodes =
    assert (Array.length nodes > 0);
    let name =
      match name with
      | Some s -> s
      | None   -> Random.int 65535 |> string_of_int
    in
    (* assumes that all the outputs are part of the same network *)
    let nn = nodes.(0).network in
    nn.nnid <- name;
    nn.outputs <- nodes;
    nn


  let get_network_name n = n.nnid

  let set_network_name n name = n.nnid <- name

  let input_shape n = (get_roots n).(0).neuron |> Neuron.get_in_shape

  let input_shapes n = Array.map (fun r -> r.neuron |> Neuron.get_in_shape) (get_roots n)

  (* collect the outputs of a given set of nodes *)
  let collect_output nodes =
    Array.map
      (fun n ->
        match n.output with
        | Some o -> o
        | None   -> failwith "Owl_neural_graph:collect_output")
      nodes


  let connect_pair prev next =
    if Array.mem prev next.prev = false
    then next.prev <- Array.append next.prev [| prev |];
    if Array.mem next prev.next = false
    then prev.next <- Array.append prev.next [| next |]


  let connect_to_parents parents child =
    (* update the child's input and output shape *)
    if Array.length parents > 0
    then (
      let out_shapes = Array.map (fun n -> n.neuron |> get_out_shape) parents in
      connect out_shapes child.neuron);
    (* connect the child to the parents *)
    Array.iter (fun p -> connect_pair p child) parents


  (* add child node to nn and connect to parents *)
  let rec add_node ?act_typ nn parents child =
    nn.size <- nn.size + 1;
    connect_to_parents parents child;
    nn.topo <- Array.append nn.topo [| child |];
    child.network <- nn;
    (* if activation is specified, recursively add_node *)
    match act_typ with
    | Some act ->
      let neuron = Activation (Activation.create act) in
      let child_of_child = make_node [||] [||] neuron None nn in
      add_node nn [| child |] child_of_child
    | None     -> child


  (* functions to interface to optimisation engine *)

  let init nn = Array.iter (fun n -> init n.neuron) nn.topo

  let reset nn = Array.iter (fun n -> reset n.neuron) nn.topo

  let mktag t nn = Array.iter (fun n -> mktag t n.neuron) nn.topo

  let mkpar nn = Array.map (fun n -> mkpar n.neuron) nn.topo

  let mkpri nn = Array.map (fun n -> mkpri n.neuron) nn.topo

  let mkadj nn = Array.map (fun n -> mkadj n.neuron) nn.topo

  let update nn us = Array.iter2 (fun n u -> update n.neuron u) nn.topo us

  let run_inputs inputs nn =
    assert (Array.(length inputs = length (get_roots nn)));
    Array.iter
      (fun n ->
        (* collect the inputs from parents' output *)
        let input =
          match n.neuron with
          | Input _ ->
            let index =
              Owl_utils.Array.index_of (Array.map (fun r -> r.name) (get_roots nn)) n.name
            in
            [| inputs.(index) |]
          | _       -> collect_output n.prev
        in
        (* process the current neuron, save output *)
        let output = run input n.neuron in
        n.output <- Some output)
      nn.topo;
    (* collect the final outputs *)
    collect_output nn.outputs


  let run x nn =
    Array.iter
      (fun n ->
        (* collect the inputs from parents' output *)
        let input =
          match n.neuron with
          | Input _ -> [| x |]
          | _       -> collect_output n.prev
        in
        (* process the current neuron, save output *)
        let output = run input n.neuron in
        n.output <- Some output)
      nn.topo;
    (* collect the final output from the tail *)
    let sink = [| nn.topo.(Array.length nn.topo - 1) |] in
    (collect_output sink).(0)


  let forward nn x =
    mktag (tag ()) nn;
    run x nn, mkpar nn


  let forward_inputs nn x =
    mktag (tag ()) nn;
    run_inputs x nn, mkpar nn


  let backward nn y =
    reverse_prop (_f 1.) y;
    mkpri nn, mkadj nn


  let copy nn =
    let nn' = make_network ~nnid:nn.nnid nn.size [||] [||] in
    (* first iteration to copy the neurons *)
    nn'.topo
      <- Array.map
           (fun node ->
             let neuron' = copy node.neuron in
             make_node ~name:node.name ~train:node.train [||] [||] neuron' None nn')
           nn.topo;
    (* second iteration to re-construct the structure and infer the shape *)
    Array.iter2
      (fun node node' ->
        node'.prev <- Array.map (fun n -> get_node nn' n.name) node.prev;
        node'.next <- Array.map (fun n -> get_node nn' n.name) node.next;
        connect_to_parents node'.prev node')
      nn.topo
      nn'.topo;
    (* set roots and outputs to finalise the structure *)
    nn'.roots <- Array.map (fun n -> get_node nn' n.name) (get_roots nn);
    nn'.outputs <- Array.map (fun n -> get_node nn' n.name) (get_outputs nn);
    nn'


  let _remove_training_nodes nn =
    let topo' =
      Owl_utils.Array.filter
        (fun n ->
          if n.train = true
          then (
            (* remove myself from my parents *)
            Array.iter
              (fun m ->
                let next' = Owl_utils.Array.filter (fun x -> x.name <> n.name) m.next in
                m.next <- next')
              n.prev;
            (* remove myself from my children *)
            Array.iter
              (fun m ->
                let prev' = Owl_utils.Array.filter (fun x -> x.name <> n.name) m.prev in
                m.prev <- prev')
              n.next;
            (* connect my parents and my children *)
            Array.iter (connect_to_parents n.prev) n.next);
          not n.train)
        nn.topo
    in
    nn.topo <- topo'


  let model nn =
    if Array.length nn.roots > 1
    then failwith "Owl_neural_graph:model Did you mean to use model_inputs?";
    let nn = copy nn in
    _remove_training_nodes nn;
    let inference x =
      match run (Arr x) nn with
      | Arr y -> y
      | _     -> failwith "Owl_neural_graph:model"
    in
    inference


  let model_inputs nn =
    let nn = copy nn in
    _remove_training_nodes nn;
    let inference inputs =
      let outputs = run_inputs (Array.map (fun x -> Arr x) inputs) nn in
      Array.map unpack_arr outputs
    in
    inference


  (* functions to create functional nodes *)

  let input ?name inputs =
    let neuron = Input (Input.create inputs) in
    let nn = make_network 0 [||] [||] in
    let n = make_node ?name [||] [||] neuron None nn in
    nn.roots <- [| n |];
    add_node nn [||] n


  let inputs ?names input_shapes =
    let names =
      match names with
      | Some x ->
        assert (Array.(length x = length input_shapes));
        Array.map (fun name -> Some name) x
      | None   -> Array.(make (length input_shapes) None)
    in
    let neurons = Array.map (fun s -> Input (Input.create s)) input_shapes in
    let nn = make_network 0 [||] [||] in
    let ns =
      Array.map2
        (fun n name -> make_node ?name [||] [||] n None nn |> add_node nn [||])
        neurons
        names
    in
    nn.roots <- ns;
    ns


  let activation ?name act_typ input_node =
    let neuron = Activation (Activation.create act_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let linear ?name ?(init_typ = Init.Standard) ?act_typ outputs input_node =
    let neuron = Linear (Linear.create outputs init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let linear_nobias ?name ?(init_typ = Init.Standard) ?act_typ outputs input_node =
    let neuron = LinearNoBias (LinearNoBias.create outputs init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let embedding ?name ?(init_typ = Init.Standard) ?act_typ in_dim out_dim input_node =
    let neuron = Embedding (Embedding.create in_dim out_dim init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let recurrent ?name ?(init_typ = Init.Standard) ~act_typ outputs hiddens input_node =
    let neuron = Recurrent (Recurrent.create hiddens outputs act_typ init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let lstm ?name ?(init_typ = Init.Tanh) cells input_node =
    let neuron = LSTM (LSTM.create cells init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let gru ?name ?(init_typ = Init.Tanh) cells input_node =
    let neuron = GRU (GRU.create cells init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let conv1d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      input_node
    =
    let neuron = Conv1D (Conv1D.create padding kernel stride init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let conv2d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      input_node
    =
    let neuron = Conv2D (Conv2D.create padding kernel stride init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let conv3d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      input_node
    =
    let neuron = Conv3D (Conv3D.create padding kernel stride init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let dilated_conv1d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      rate
      input_node
    =
    let neuron =
      DilatedConv1D (DilatedConv1D.create padding kernel stride rate init_typ)
    in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let dilated_conv2d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      rate
      input_node
    =
    let neuron =
      DilatedConv2D (DilatedConv2D.create padding kernel stride rate init_typ)
    in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let dilated_conv3d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      rate
      input_node
    =
    let neuron =
      DilatedConv3D (DilatedConv3D.create padding kernel stride rate init_typ)
    in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let transpose_conv1d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      input_node
    =
    let neuron =
      TransposeConv1D (TransposeConv1D.create padding kernel stride init_typ)
    in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let transpose_conv2d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      input_node
    =
    let neuron =
      TransposeConv2D (TransposeConv2D.create padding kernel stride init_typ)
    in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let transpose_conv3d
      ?name
      ?(padding = SAME)
      ?(init_typ = Init.Tanh)
      ?act_typ
      kernel
      stride
      input_node
    =
    let neuron =
      TransposeConv3D (TransposeConv3D.create padding kernel stride init_typ)
    in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let fully_connected ?name ?(init_typ = Init.Standard) ?act_typ outputs input_node =
    let neuron = FullyConnected (FullyConnected.create outputs init_typ) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let max_pool1d ?name ?(padding = SAME) ?act_typ kernel stride input_node =
    let neuron = MaxPool1D (MaxPool1D.create padding kernel stride) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let max_pool2d ?name ?(padding = SAME) ?act_typ kernel stride input_node =
    let neuron = MaxPool2D (MaxPool2D.create padding kernel stride) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let avg_pool1d ?name ?(padding = SAME) ?act_typ kernel stride input_node =
    let neuron = AvgPool1D (AvgPool1D.create padding kernel stride) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let avg_pool2d ?name ?(padding = SAME) ?act_typ kernel stride input_node =
    let neuron = AvgPool2D (AvgPool2D.create padding kernel stride) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let global_max_pool1d ?name ?act_typ input_node =
    let neuron = GlobalMaxPool1D (GlobalMaxPool1D.create ()) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let global_max_pool2d ?name ?act_typ input_node =
    let neuron = GlobalMaxPool2D (GlobalMaxPool2D.create ()) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let global_avg_pool1d ?name ?act_typ input_node =
    let neuron = GlobalAvgPool1D (GlobalAvgPool1D.create ()) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let global_avg_pool2d ?name ?act_typ input_node =
    let neuron = GlobalAvgPool2D (GlobalAvgPool2D.create ()) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let upsampling2d ?name ?act_typ size input_node =
    let neuron = UpSampling2D (UpSampling2D.create size) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let padding2d ?name ?act_typ padding input_node =
    let neuron = Padding2D (Padding2D.create padding) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let dropout ?name rate input_node =
    let neuron = Dropout (Dropout.create rate) in
    let nn = get_network input_node in
    let n = make_node ?name ~train:true [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let gaussian_noise ?name sigma input_node =
    let neuron = GaussianNoise (GaussianNoise.create sigma) in
    let nn = get_network input_node in
    let n = make_node ?name ~train:true [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let gaussian_dropout ?name rate input_node =
    let neuron = GaussianDropout (GaussianDropout.create rate) in
    let nn = get_network input_node in
    let n = make_node ?name ~train:true [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let alpha_dropout ?name rate input_node =
    let neuron = AlphaDropout (AlphaDropout.create rate) in
    let nn = get_network input_node in
    let n = make_node ?name ~train:true [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let normalisation ?name ?(axis = -1) ?training ?decay ?mu ?var input_node =
    let neuron = Normalisation (Normalisation.create ?training ?decay ?mu ?var axis) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let reshape ?name outputs input_node =
    let neuron = Reshape (Reshape.create outputs) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let flatten ?name input_node =
    let neuron = Flatten (Flatten.create ()) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let slice ?name slice input_node =
    let neuron = Slice (Slice.create slice) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node nn [| input_node |] n


  let lambda ?name ?act_typ ?out_shape lambda input_node =
    let neuron = Lambda (Lambda.create ?out_shape lambda) in
    let nn = get_network input_node in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn [| input_node |] n


  let lambda_array ?name ?act_typ out_shape lambda input_node =
    let neuron = LambdaArray (LambdaArray.create out_shape lambda) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  let add ?name ?act_typ input_node =
    let neuron = Add (Add.create ()) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  let mul ?name ?act_typ input_node =
    let neuron = Mul (Mul.create ()) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  let dot ?name ?act_typ input_node =
    let neuron = Dot (Dot.create ()) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  let max ?name ?act_typ input_node =
    let neuron = Max (Max.create ()) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  let average ?name ?act_typ input_node =
    let neuron = Average (Average.create ()) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  let concatenate ?name ?act_typ axis input_node =
    let neuron = Concatenate (Concatenate.create axis) in
    let nn = get_network input_node.(0) in
    let n = make_node ?name [||] [||] neuron None nn in
    add_node ?act_typ nn input_node n


  (* I/O functions *)

  let to_string nn =
    let s = ref (nn.nnid ^ "\n\n") in
    Array.iter
      (fun n ->
        let prev =
          Array.map (fun n -> n.name) n.prev |> Owl_utils_array.to_string (fun s -> s)
        in
        let next =
          Array.map (fun n -> n.name) n.next |> Owl_utils_array.to_string (fun s -> s)
        in
        s
          := !s
             ^ Printf.sprintf "\x1b[31m[ Node %s ]:\x1b[0m\n" n.name
             ^ Printf.sprintf "%s" (to_string n.neuron)
             ^ Printf.sprintf "    prev:[%s] next:[%s]\n\n" prev next)
      nn.topo;
    !s


  let pp_network formatter nn =
    Format.open_box 0;
    Format.fprintf formatter "%s" (to_string nn);
    Format.close_box ()


  let print nn = pp_network Format.std_formatter nn

  let save ?(unsafe = false) nn f =
    if unsafe = true
    then (
      Owl_log.warn
        "Unsafely saved network can only be loaded back in exactly the same version of \
         OCaml and Owl.";
      Owl_io.marshal_to_file ~flags:[ Marshal.Closures ] (copy nn) f)
    else Owl_io.marshal_to_file (copy nn) f


  let load f : network = Owl_io.marshal_from_file f

  let save_weights nn f =
    let h = Hashtbl.create nn.size in
    Array.iter
      (fun n ->
        let ws = Neuron.save_weights n.neuron in
        Hashtbl.add h n.name ws)
      nn.topo;
    Owl_io.marshal_to_file h f


  let load_weights nn f =
    let h = Owl_io.marshal_from_file f in
    Array.iter
      (fun n ->
        let ws = Hashtbl.find h n.name in
        Neuron.load_weights n.neuron ws)
      nn.topo


  let make_subnetwork ?(copy = true) ?(make_inputs = [||]) nn output_names =
    let subnn = make_network 0 [||] [||] in
    let in_nodes = ref [] in
    (* collect neurons belonging to subnetwork *)
    let rec collect_subnn_nodes n acc =
      if List.exists (fun in_acc -> in_acc.name = n.name) acc
      then acc
      else if Array.mem n.name make_inputs
      then (
        let shape = get_out_shape n.neuron in
        let in_neur = Input (Input.create shape) in
        let new_in = make_node ~name:n.name [||] [||] in_neur None subnn in
        in_nodes := new_in :: !in_nodes;
        new_in :: acc)
      else (
        let neur = if copy then Neuron.copy n.neuron else n.neuron in
        let new_node = make_node ~name:n.name ~train:n.train [||] [||] neur None subnn in
        match neur with
        | Input _ ->
          in_nodes := new_node :: !in_nodes;
          new_node :: acc
        | _       ->
          let acc = new_node :: acc in
          Array.fold_left (fun a prev -> collect_subnn_nodes prev a) acc n.prev)
    in
    let new_nodes =
      Array.fold_left
        (fun acc name -> collect_subnn_nodes (get_node nn name) acc)
        []
        output_names
    in
    (* sorts the new topology *)
    let new_topo =
      Array.fold_left
        (fun acc n ->
          match List.find_opt (fun n' -> n'.name = n.name) new_nodes with
          | Some n' -> n' :: acc
          | None    -> acc)
        []
        nn.topo
      |> List.rev
      |> Array.of_list
    in
    subnn.topo <- new_topo;
    (* re-construct network structure *)
    Array.iter
      (fun node' ->
        let node = get_node nn node'.name in
        if not (List.memq node' !in_nodes)
        then node'.prev <- Array.map (fun n -> get_node subnn n.name) node.prev;
        if not (Array.mem node.name output_names)
        then (
          (* only process nodes that are part of the subnetwork *)
          let next =
            Owl_utils_array.filter
              (fun n -> Array.exists (fun n' -> n'.name = n.name) subnn.topo)
              node.next
          in
          (* With custom input nodes, next could contain an input node. *)
          node'.next <- Array.map (fun n -> get_node subnn n.name) next);
        connect_to_parents node'.prev node')
      subnn.topo;
    (* TODO: Warn if not all names in in_names were used? *)
    subnn.roots <- Array.of_list !in_nodes;
    subnn.outputs <- Array.map (fun name -> get_node subnn name) output_names;
    subnn


  (* training functions *)

  (* generic minimisation functions
     forward: function to run the forward pass
     backward: function to run the backward pass
     update: function to update the weights according to the gradient
     save: function to save the model for checkpoint
  *)
  let train_generic ?state ?params ?(init_model = true) nn x y =
    if init_model = true then init nn;
    let f = forward nn in
    let b = backward nn in
    let u = update nn in
    let s = save nn in
    let p =
      match params with
      | Some p -> p
      | None   -> Optimise.Params.default ()
    in
    Optimise.minimise_network ?state p f b u s x y


  let train ?state ?params ?init_model nn x y =
    train_generic ?state ?params ?init_model nn (Arr x) (Arr y)
end

(* Make functor ends *)