Tiny cleanup of nn_blocks.ml

lukstafi · lukstafi · commit a2c3c372c731 · 2025-08-28T11:27:15.000+02:00
diff --git a/lib/nn_blocks.ml b/lib/nn_blocks.ml
@@ -1,15 +1,12 @@
-(** Prior to OCANNL 0.5, this module is just a placeholder hinting at an intended design pattern for
-    model components. *)
-
 open! Base
 module TDSL = Operation.TDSL
 module NTDSL = Operation.NTDSL
 
 let%op mlp_layer ~label ~hid_dim () x = relu (({ w = uniform () } * x) + { b = 0.; o = [ hid_dim ] })
 
-let mlp ~hid_dims =
+let mlp ~label ~hid_dims () =
   let layers =
     List.mapi hid_dims ~f:(fun i hid_dim ->
-        mlp_layer ~label:[ "L" ^ Int.to_string i ] ~hid_dim ())
+        mlp_layer ~label:(("L" ^ Int.to_string i) :: label) ~hid_dim ())
   in
   fun x -> List.fold layers ~init:x ~f:(fun x layer -> layer x)
