Merge pull request #378 from ahrefs/fix-label-parameter-type

lukstafi · web-flow · commit c8da5c469734 · 2025-08-27T22:56:55.000+02:00
Fix ~label parameter type in %op syntax extension
diff --git a/lib/nn_blocks.ml b/lib/nn_blocks.ml
@@ -5,15 +5,11 @@ open! Base
 module TDSL = Operation.TDSL
 module NTDSL = Operation.NTDSL
 
-type mlp_layer_config = { label : string list; hid_dim : int }
+let%op mlp_layer ~label ~hid_dim () x = relu (({ w = uniform () } * x) + { b = 0.; o = [ hid_dim ] })
 
-let%op mlp_layer ~config x = relu (({ w = uniform () } * x) + { b = 0.; o = [ config.hid_dim ] })
-
-type mlp_config = { label : string list; hid_dims : int list }
-
-let mlp ~config =
+let mlp ~hid_dims =
   let layers =
-    List.mapi config.hid_dims ~f:(fun i hid_dim ->
-        mlp_layer ~config:{ label = [ "L" ^ Int.to_string i ]; hid_dim })
+    List.mapi hid_dims ~f:(fun i hid_dim ->
+        mlp_layer ~label:[ "L" ^ Int.to_string i ] ~hid_dim ())
   in
   fun x -> List.fold layers ~init:x ~f:(fun x layer -> layer x)
diff --git a/lib/ppx_op.ml b/lib/ppx_op.ml
@@ -3,8 +3,10 @@ open Ppxlib
 open Ppx_arrayjit.Ppx_helper
 open Ppx_shared
 
-let make_p ~has_config ~loc ?value ?values ?param_init ~extra_args name =
-  let more_label = if has_config then [%expr Some config.label] else [%expr None] in
+let make_p ~opt_label ~loc ?value ?values ?param_init ~extra_args name =
+  let more_label = match opt_label with 
+    | Some (_label_name, label_pat) -> [%expr Some [%e pat2expr label_pat]]
+    | None -> [%expr None] in
   let value = match value with Some c -> [%expr Some [%e c]] | None -> [%expr None] in
   let values = match values with Some c -> [%expr Some [%e c]] | None -> [%expr None] in
   let param_init =
@@ -36,13 +38,13 @@ let make_p ~has_config ~loc ?value ?values ?param_init ~extra_args name =
   in
   [%expr [%e with_extra_args] ()]
 
-let make_vb ~has_config ?value ?param_init ~extra_args ~loc name =
+let make_vb ~opt_label ?value ?param_init ~extra_args ~loc name =
   let pat = Ast_helper.Pat.var ~loc:name.loc name in
-  let v = make_p ~has_config ~loc ?value ?param_init ~extra_args name in
+  let v = make_p ~opt_label ~loc ?value ?param_init ~extra_args name in
   let vb = Ast_helper.Vb.mk ~loc pat v in
   (pat, vb)
 
-let make_vb_nd ~has_config ~init_nd ~extra_args ~loc name =
+let make_vb_nd ~opt_label ~init_nd ~extra_args ~loc name =
   let pat = Ast_helper.Pat.var ~loc:name.loc name in
   let values, batch_dims, output_dims, input_dims = ndarray_constant init_nd in
   let v =
@@ -59,7 +61,7 @@ let make_vb_nd ~has_config ~init_nd ~extra_args ~loc name =
         :: ({ txt = Lident "output_dims"; loc }, output_dims_expr)
         :: extra_args
       in
-      make_p ~has_config ~loc ~values ~extra_args name
+      make_p ~opt_label ~loc ~values ~extra_args name
   in
   let vb = Ast_helper.Vb.mk ~loc pat v in
   (pat, vb)
@@ -80,9 +82,9 @@ let lift_config_vb ~loop ~num_configs ?label ~expr1 ~c_expr arg_exprs =
     | [ e2; e3 ] -> [%expr [%e pat2expr pat] [%e e2] [%e e3]]
     | _ -> assert false )
 
-let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
+let rec translate ~num_configs ~is_toplevel ~opt_label ?label expr =
   let loc = expr.pexp_loc in
-  let loop = translate ~num_configs ~is_toplevel:false ~has_config in
+  let loop = translate ~num_configs ~is_toplevel:false ~opt_label in
   match expr with
   | { pexp_desc = Pexp_constant (Pconst_float _); _ } ->
       (no_vbs, [%expr TDSL.number ?label:[%e opt_expr ~loc label] [%e expr]])
@@ -141,7 +143,7 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
       match first_label.txt with
       | Lident tensor_name ->
           let name = { loc = first_label.loc; txt = tensor_name } in
-          let pat, vb = make_vb ~has_config ~value ~extra_args ~loc name in
+          let pat, vb = make_vb ~opt_label ~value ~extra_args ~loc name in
           (Map.singleton (module String) tensor_name vb, pat2expr pat)
       | _ ->
           ( no_vbs,
@@ -160,7 +162,7 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
       | Lident tensor_name ->
           let value = [%expr Float.of_int [%e int_val]] in
           let name = { loc = first_label.loc; txt = tensor_name } in
-          let pat, vb = make_vb ~has_config ~value ~extra_args ~loc name in
+          let pat, vb = make_vb ~opt_label ~value ~extra_args ~loc name in
           (Map.singleton (module String) tensor_name vb, pat2expr pat)
       | _ ->
           ( no_vbs,
@@ -181,7 +183,7 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
       match first_label.txt with
       | Lident tensor_name ->
           let name = { loc = first_label.loc; txt = tensor_name } in
-          let pat, vb = make_vb_nd ~has_config ~init_nd ~extra_args ~loc name in
+          let pat, vb = make_vb_nd ~opt_label ~init_nd ~extra_args ~loc name in
           (* Note: expect a type error if batch_dims exist or extra_args modify the shape *)
           (Map.singleton (module String) tensor_name vb, pat2expr pat)
       | _ ->
@@ -204,7 +206,7 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
                 (vbs, Some e)
           in
           let name = { loc = first_label.loc; txt = tensor_name } in
-          let pat, vb = make_vb ~has_config ?param_init ~extra_args ~loc name in
+          let pat, vb = make_vb ~opt_label ?param_init ~extra_args ~loc name in
           (* Combine with any bindings from the initialization *)
           let all_vbs = Map.add_exn init_vbs ~key:tensor_name ~data:vb in
           (all_vbs, pat2expr pat)
@@ -258,29 +260,44 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
       let vbs1, e1 = loop ?label expr1 in
       let vbs2, e2 = loop expr2 in
       (reduce_vbss [ vbs1; vbs2 ], [%expr [%e e1] [%e e2]])
-  | {
-   pexp_desc =
-     Pexp_function
-       ( ({ pparam_desc = Pparam_val (Labelled "config", c_e, c_pat); _ } as arg) :: args,
-         constr,
-         body );
-   _;
-  } ->
-      let vbs, body =
-        translate ~num_configs ~is_toplevel:true ~has_config:true ?label
-          { expr with pexp_desc = Pexp_function (args, constr, body) }
+  | { pexp_desc = Pexp_function (args, constr, body); _ } when is_toplevel -> (
+      (* Check if there's a unit parameter or a labeled parameter with label "label" *)
+      let rec find_unit_pos idx = function
+        | [] -> None
+        | { pparam_desc = Pparam_val (Nolabel, _, pat); _ } :: _ 
+          when match pat.ppat_desc with
+               | Ppat_construct ({ txt = Lident "()"; _ }, None) -> true
+               | _ -> false ->
+            Some idx
+        | _ :: rest -> find_unit_pos (idx + 1) rest
       in
-      let body = let_opt ~loc vbs body in
-      ( no_vbs,
-        {
-          expr with
-          pexp_desc =
-            Pexp_function
-              ( [ { arg with pparam_desc = Pparam_val (Labelled "config", c_e, c_pat) } ],
-                constr,
-                Pfunction_body body );
-        } )
-  | { pexp_desc = Pexp_function (args, constr, body); _ } when is_toplevel ->
+      let rec find_label_param = function
+        | [] -> None
+        | { pparam_desc = Pparam_val (Labelled "label", _, pat); _ } :: _ -> Some ("label", pat)
+        | _ :: rest -> find_label_param rest
+      in
+      match find_unit_pos 0 args with
+      | Some unit_idx ->
+          (* Split args at unit parameter *)
+          let before_unit, unit_and_after = List.split_n args unit_idx in
+          let unit_param, after_unit = match unit_and_after with
+            | unit :: rest -> (unit, rest)
+            | [] -> failwith "Internal error: unit_and_after should not be empty" in
+          let opt_label = find_label_param before_unit in
+          let vbs, inner_body =
+            translate ~num_configs ~is_toplevel:false ~opt_label ?label
+              { expr with pexp_desc = Pexp_function (after_unit, constr, body) }
+          in
+          let inner_body = let_opt ~loc vbs inner_body in
+          (* The inner_body already has after_unit parameters processed, so use it directly *)
+          let new_body = inner_body in
+          ( no_vbs,
+            if List.is_empty before_unit then
+              { expr with pexp_desc = Pexp_function ([unit_param], constr, Pfunction_body new_body) }
+            else
+              { expr with pexp_desc = Pexp_function (before_unit @ [unit_param], constr, Pfunction_body new_body) } )
+      | None ->
+      (* No unit parameter, normal processing *)
       let labels =
         Option.to_list label
         @ List.filter_map args ~f:(function
@@ -323,7 +340,7 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
                 ~f:(fun acc vbs -> Map.merge_disjoint_exn acc vbs),
               Pfunction_cases (cases, loc, attrs) )
       in
-      (vbs, { expr with pexp_desc = Pexp_function (args, constr, body) })
+      (vbs, { expr with pexp_desc = Pexp_function (args, constr, body) }) )
   | { pexp_desc = Pexp_function (args, constr, body); _ } ->
       let vbs, body =
         match body with
@@ -422,7 +439,7 @@ let rec translate ~num_configs ~is_toplevel ~has_config ?label expr =
 
 let translate ?ident_label expr =
   let vbs, expr =
-    translate ~num_configs:(ref 0) ~is_toplevel:true ~has_config:false
+    translate ~num_configs:(ref 0) ~is_toplevel:true ~opt_label:None
       ~label:(opt_pat2string_list ~loc:expr.pexp_loc ident_label)
       expr
   in
diff --git a/lib/syntax_extensions.md b/lib/syntax_extensions.md
@@ -307,12 +307,10 @@ Inline declarations can also be used outside of assignments for creating non-dif
   in
 ```
 
-For `%op`, the declaration is allowed anywhere. If there is a `~config` function parameter used inside the extension scope, for example as `fun ~config ... -> ...` or a more specific example `let%op mlp ~config x = ...`, the scope of an inline-declared tensor is no longer the full scope of the extension point. Instead, the tensor is defined right underneath the introduction of the `~config` parameter: `fun ~config -> let <definitions of the inline-declared tensors> in ...`. The config value passed to the generated code must be a record with at least a field `label : string list`. The inline-declared tensor that's defined under a `~config` parameter is defined as `TDSL.param ~more_label:config.label ...` Example showing two param tensors declared inline, including `config.label` in their labels:
+For `%op`, the declaration is allowed anywhere. If there is a unit `()` parameter in the function, the scope of inline-declared tensors is delimited at that parameter. The tensors are defined right after the unit parameter. If there is a labeled parameter with label `label` before the unit parameter (e.g., `~label`), the inline-declared tensors will use that parameter (which should be of type `string list`) to enrich their labels. Example showing two param tensors declared inline, with scope delimited by `()` and labels enriched by the `label` parameter:
 
 ```ocaml
-type mlp_layer_config = { label : string list; hid_dim : int }
-
-let%op mlp_layer ~config x = relu ({ w } * x + { b; o = [ config.hid_dim ] })
+let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] })
 ```
 
 ## Using OCANNL's generalized einsum notation
@@ -418,13 +416,15 @@ This syntax used to be very important, because comments in assignments are used
 
 ### Name from binding
 
-When an extension point is applied to a let-binding, e.g. `let%op mlp_layer ~config x = relu ({ w } * x + { b; o = [ config.hid_dim ] })`, it uses the name of the binding (`mlp_layer` in the example) for the label of the primary tensor created by the extension, if any. This is why the resulting layer tensor in the example has its label starting with `"mlp_layer"`. If the extension is over a semicolon-separated sequence of expressions, the primary tensor can only be in the last component of the sequence, other syntax constructs are handled analogously.
+When an extension point is applied to a let-binding, e.g. `let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] })`, it uses the name of the binding (`mlp_layer` in the example) for the label of the primary tensor created by the extension, if any. This is why the resulting layer tensor in the example has its label starting with `"mlp_layer"`. If the extension is over a semicolon-separated sequence of expressions, the primary tensor can only be in the last component of the sequence, other syntax constructs are handled analogously.
+
+The example `let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] })` also illustrates providing additional string list to populate the label of the tensor: `label` must be of type `string list`.
 
 ### Label from function argument
 
 The resulting (primary) tensor's label will also have incorporated the label of the input argument, if any. In our example, the resulting `mlp_layer` tensor will also include the label of the actually applied `x`.
 
-Note that we do not include `config.label`, even if `config` is available, because the actually applied input argument will typically have more specific information.
+Note that we do not include separate config labels, because the actually applied input argument will typically have more specific information.
 
 ### Configuring inline declarations: inline output dimensions, initial values
 
@@ -449,59 +449,51 @@ A very simple example from [micrograd_demo: Micrograd README basic example](test
 
 ### Lifting of the applications of config arguments: if an error, refactor your code
 
-If you recall, inline declared param tensors get lifted out of functions except for the function `fun ~config ->`, where they get defined. Our example `let%op mlp_layer ~config x = relu ({ w } * x + { b; o = [ config.hid_dim ] })` translates as:
+If you recall, inline declared param tensors get lifted out of functions to be defined at the point of a unit `()` parameter. Our example `let%op mlp_layer ~label ~hid_dim () x = relu ({ w } * x + { b; o = [ hid_dim ] })` translates as:
 
 ```ocaml
-let mlp_layer ~config =
-  let w = TDSL.param ~more_label:config.label "w" () 
-  and b = TDSL.param ~more_label:config.label ~output_dims:[ config.hid_dim ] "b" () in
+let mlp_layer ~label ~hid_dim () =
+  let w = TDSL.param ~more_label:label "w" () 
+  and b = TDSL.param ~more_label:label ~output_dims:[ hid_dim ] "b" () in
   fun x -> TDSL.O.(relu (w * x + b))
 ```
 
 For this to work properly, when employing such network blocks, their params also need to be introduced at the right moment. Therefore, the `%op` syntax ensures that this example:
 
 ```ocaml
-type tlp_config = { label : string list; dim1 : int; dim2 : int; dim3 : int }
-
-let%op three_layer_perceptron ~config x =
-  mlp_layer ~config:{ label = [ "L3" ]; hid_dim = config.dim3 }
-    (mlp_layer ~config:{ label = [ "L2" ]; hid_dim = config.dim2 }
-       (mlp_layer ~config:{ label = [ "L1" ]; hid_dim = config.dim1 } x))
+let%op three_layer_perceptron ~label ~dim1 ~dim2 ~dim3 () x =
+  mlp_layer ~label:[ "L3" ] ~hid_dim:dim3 ()
+    (mlp_layer ~label:[ "L2" ] ~hid_dim:dim2 ()
+       (mlp_layer ~label:[ "L1" ] ~hid_dim:dim1 () x))
 ```
 
 gets expanded to:
 
 ```ocaml
-type tlp_config = { label : string list; dim1 : int; dim2 : int; dim3 : int }
-
-let three_layer_perceptron ~config =
-  let config_block__1 = mlp_layer ~config:{ label = [ "L3" ]; hid_dim = config.dim3 }
-  and config_block__2 = mlp_layer ~config:{ label = [ "L2" ]; hid_dim = config.dim2 }
-  and config_block__3 = mlp_layer ~config:{ label = [ "L1" ]; hid_dim = config.dim1 } in
+let three_layer_perceptron ~label ~dim1 ~dim2 ~dim3 () =
+  let config_block__1 = mlp_layer ~label:[ "L3" ] ~hid_dim:dim3 ()
+  and config_block__2 = mlp_layer ~label:[ "L2" ] ~hid_dim:dim2 ()
+  and config_block__3 = mlp_layer ~label:[ "L1" ] ~hid_dim:dim1 () in
   fun x -> config_block__1 (config_block__2 (config_block__3 x))
 ```
 
 However, this raises a concern for more complex situations. Consider this code that fails to compile:
 
 ```ocaml
-type mlp_config = { label : string list; hid_dims : int list }
-
-let%op mlp ~config x =
-  List.foldi config.hid_dims ~init:x ~f:(fun i x hid_dim ->
-      mlp_layer ~config:{ label = [ "L" ^ Int.to_string i ]; hid_dim } x)
+let%op mlp ~label ~hid_dims () x =
+  List.foldi hid_dims ~init:x ~f:(fun i x hid_dim ->
+      mlp_layer ~label:[ "L" ^ Int.to_string i ] ~hid_dim () x)
 ```
 
 The attempted lifting breaks because of the escaping variables `i` and `hid_dim`. This reminds us to rewrite the example, ensuring the proper introduction of params:
 
 ```ocaml
-type mlp_config = { label : string list; hid_dims : int list }
-
-let mlp ~config =
+let mlp ~label ~hid_dims =
   let layers =
-    List.mapi config.hid_dims ~f:(fun i hid_dim ->
-        mlp_layer ~config:{ label = [ "L" ^ Int.to_string i ]; hid_dim })
+    List.mapi hid_dims ~f:(fun i hid_dim ->
+        mlp_layer ~label:[ "L" ^ Int.to_string i ] ~hid_dim ())
   in
-  fun x -> List.fold layers ~init:x ~f:(fun x layer -> layer x)
+  fun () x -> List.fold layers ~init:x ~f:(fun x layer -> layer x)
 ```
 
 Unfortunately, we need to be mindful to introduce params at the right times.
diff --git a/test/ppx/test_ppx_op.ml b/test/ppx/test_ppx_op.ml
@@ -21,21 +21,17 @@ let z3 =
 
 let () = ignore (y0, y1, y2, a, b, y, z, z2, z3)
 
-type mlp_layer_config = { label : string list; hid_dim : int }
-
-let%op mlp_layer ~config x = relu (({ w } * x) + { b; o = [ config.hid_dim ] })
+let%op mlp_layer ~label ~hid_dim () x = relu (({ w } * x) + { b; o = [ hid_dim ] })
 
 let%op _use_layer x =
-  mlp_layer ~config:{ label = [ "L" ]; hid_dim = 3 }
-    (mlp_layer ~config:{ label = [ "L2" ]; hid_dim = 3 } x)
-
-let%op _config_layer ~config:_ x = mlp_layer ~config:{ label = [ "L" ]; hid_dim = 3 } x
+  mlp_layer ~label:[ "L" ] ~hid_dim:3 ()
+    (mlp_layer ~label:[ "L2" ] ~hid_dim:3 () x)
 
-type tlp_config = { label : string list; dim1 : int; dim2 : int; dim3 : int }
+let%op _config_layer ~config:_ x = mlp_layer ~label:[ "L" ] ~hid_dim:3 () x
 
-let%op _three_layer_perceptron ~(config : tlp_config) x =
+let%op _three_layer_perceptron ~label ~dim1 ~dim2 ~dim3 () x =
   mlp_layer
-    ~config:{ label = "L3" :: config.label; hid_dim = config.dim3 }
+    ~label:(label @ [ "L3" ]) ~hid_dim:dim3 ()
     (mlp_layer
-       ~config:{ label = "L2" :: config.label; hid_dim = config.dim2 }
-       (mlp_layer ~config:{ label = "L1" :: config.label; hid_dim = config.dim1 } x))
+       ~label:(label @ [ "L2" ]) ~hid_dim:dim2 ()
+       (mlp_layer ~label:(label @ [ "L1" ]) ~hid_dim:dim1 () x))
diff --git a/test/ppx/test_ppx_op_expected.ml b/test/ppx/test_ppx_op_expected.ml
