The Nn_blocks placeholder hinting at intended design of model components

lukstafi · lukstafi · commit efb5c81b565d · 2024-09-12T21:50:52.000+02:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -7,6 +7,7 @@
   - We check FP16 constants for overflow.
   - We output half precision specific code from the CUDA backend.
 - Finally proper support for mixed precision! Lazy precision defaults and delayed precision setting via `Tnode.update_prec`.
+- A placeholder `nn_blocks.ml` hinting at an intended design pattern for model components.
 
 ### Changed
 
diff --git a/lib/dune b/lib/dune
@@ -16,7 +16,7 @@
   arrayjit)
  (preprocess
   (pps ppx_jane ppx_ocannl ppx_minidebug))
- (modules PrintBox_utils row shape tensor operation train)
+ (modules PrintBox_utils row shape tensor operation train nn_blocks)
  (modes byte native)
  (c_library_flags -pthread))
 
diff --git a/lib/nn_blocks.ml b/lib/nn_blocks.ml
@@ -0,0 +1,19 @@
+(** 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
+
+type mlp_layer_config = { label : string list; hid_dim : int }
+
+let%op mlp_layer ~config x = ?/(("w" * x) + "b" config.hid_dim)
+
+type mlp_config = { label : string list; hid_dims : int list }
+
+let mlp ~config =
+  let layers =
+    List.mapi config.hid_dims ~f:(fun i hid_dim ->
+        mlp_layer ~config:{ label = [ "L" ^ Int.to_string i ]; hid_dim })
+  in
+  fun x -> List.fold layers ~init:x ~f:(fun x layer -> layer x)
