Enhance %cd syntax to allow inline tensor declarations in standalone expressions. Adjust related errors, comments and documentation for consistency.

lukstafi · lukstafi · commit db42dc47f0da · 2025-07-01T14:59:01.000+02:00
diff --git a/lib/ppx_cd.ml b/lib/ppx_cd.ml
@@ -111,6 +111,7 @@ let assignment ~punned ~lhs ~rhses body =
     | _ -> (no_vbs, body)
   in
   let body =
+    (* Note: this is not a binding from an inline declaration, it's a temporary binding. *)
     if Option.is_some lhs.vb then
       Ast_builder.Default.pexp_extension ~loc
       @@ Location.error_extensionf ~loc
@@ -200,7 +201,6 @@ let empty_comp ~loc = [%expr { asgns = Ir.Assignments.Noop; embedded_nodes = [%e
 
 let setup_array ~punned ~bad_pun_hints ~is_lhs
     { typ = filler_typ; slot; expr = filler; vbs; array_opt_of_code } =
-  assert (Map.is_empty vbs);
   let loc = filler.pexp_loc in
   let opt_buffer tn =
     if is_lhs then [%expr Some [%e tn]] else [%expr Some (Ir.Assignments.Node [%e tn])]
@@ -220,17 +220,18 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
       pun_hint_tnode;
     }
   in
-  match filler_typ with
-  | No_grad_tensor_intro _ when not is_lhs ->
+  match (Map.is_empty vbs, filler_typ) with
+  | (false, _ | _, No_grad_tensor_intro _) when not is_lhs ->
       {
         default_setup with
         array_opt =
           Ast_builder.Default.pexp_extension ~loc
           @@ Location.error_extensionf ~loc
-               "ppx_ocannl %%cd: punning is only allowed in the assigned-to position";
+               "ppx_ocannl %%cd: inline tensor declarations are not allowed in assignment \
+                right-hand side, to prevent over-use in locations with less label information";
       }
-  | (Tensor | Unknown) when match filler with { pexp_desc = Pexp_ident _; _ } -> true | _ -> false
-    ->
+  | _, (Tensor | Unknown)
+    when match filler with { pexp_desc = Pexp_ident _; _ } -> true | _ -> false ->
       let t = filler in
       let fwd_code_or_noop =
         Some
@@ -241,7 +242,7 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
             else [%e empty_comp ~loc]]
       in
       { default_setup with fwd_code_or_noop; tensor = Some t }
-  | Value_of_tensor ({ pexp_desc = Pexp_ident _; _ } as t) ->
+  | _, Value_of_tensor ({ pexp_desc = Pexp_ident _; _ } as t) ->
       let fwd_code_or_noop =
         Some
           [%expr
@@ -256,7 +257,7 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
         array_opt = opt_buffer [%expr [%e t].Tensor.value];
         tensor = Some t;
       }
-  | Value_of_tensor t ->
+  | _, Value_of_tensor t ->
       {
         default_setup with
         array_opt =
@@ -266,7 +267,7 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
                 identifier";
         tensor = Some t;
       }
-  | Tensor | Unknown ->
+  | _, (Tensor | Unknown) ->
       (* Need to bind the expression computing the tensor so we don't recompute it. *)
       let v =
         match slot with
@@ -293,7 +294,7 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
         array_opt = opt_buffer [%expr [%e t].Tensor.value];
         tensor = Some t;
       }
-  | No_grad_tensor_intro _ ->
+  | _, No_grad_tensor_intro _ ->
       (* Inline tensors are guaranteed to be leaf tensors, so they don't have forward code, but they
          are embedded. *)
       let fwd_code_or_noop =
@@ -306,7 +307,7 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
               }]
       in
       { default_setup with fwd_code_or_noop; tensor = Some filler }
-  | Code when Option.is_none array_opt_of_code ->
+  | _, Code when Option.is_none array_opt_of_code ->
       {
         default_setup with
         fwd_code_or_noop = Some filler;
@@ -315,16 +316,16 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
           @@ Location.error_extensionf ~loc
                "ppx_ocannl %%cd: could not determine a lead array of provided code";
       }
-  | Code ->
+  | _, Code ->
       {
         default_setup with
         fwd_code_or_noop = Some filler;
         array_opt = buffer (Option.value_exn array_opt_of_code);
       }
-  | Array -> { default_setup with array_opt = opt_buffer filler }
-  | Grad_of_tensor ({ pexp_desc = Pexp_ident _; _ } as t) ->
+  | _, Array -> { default_setup with array_opt = opt_buffer filler }
+  | _, Grad_of_tensor ({ pexp_desc = Pexp_ident _; _ } as t) ->
       { default_setup with array_opt = buffer filler; tensor = Some t }
-  | Grad_of_tensor t ->
+  | _, Grad_of_tensor t ->
       {
         default_setup with
         array_opt =
@@ -334,16 +335,16 @@ let setup_array ~punned ~bad_pun_hints ~is_lhs
                 identifier";
         tensor = Some t;
       }
-  | (Merge_value _ | Merge_grad _) when is_lhs ->
+  | _, (Merge_value _ | Merge_grad _) when is_lhs ->
       {
         default_setup with
         array_opt =
           Ast_builder.Default.pexp_extension ~loc
           @@ Location.error_extensionf ~loc "ppx_ocannl %%cd: merge buffers cannot be assigned to";
       }
-  | Merge_value t ->
+  | _, Merge_value t ->
       { default_setup with array_opt = [%expr Some (Merge_buffer [%e filler])]; tensor = Some t }
-  | Merge_grad t ->
+  | _, Merge_grad t ->
       {
         default_setup with
         array_opt = [%expr Option.map [%e filler] ~f:(fun tn -> Ir.Assignments.Merge_buffer tn)];
@@ -440,7 +441,7 @@ let translate (expr : expression) : result =
                @@ Location.error_extensionf ~loc
                     "ppx_ocannl %%cd: expected a ternary operator, one of: where, fma" ))
     in
-    (* FIXME: collapse these (code reuse) *)
+    (* TODO: collapse these (code reuse) *)
     let process_assign_ternop ~accu_op ~lhs ~tern_op ~rhs1 ~rhs2 ~rhs3 ?projections ~proj_in_scope
         () =
       let initialize_neutral, accu_op = assignment_op accu_op in
@@ -707,10 +708,16 @@ let translate (expr : expression) : result =
           slot = Scalar;
         }
     | { pexp_desc = Pexp_constant (Pconst_string (name, str_loc, _)); _ } ->
+        (* TODO: consider passing toplevel binding name as a hint label *)
+        let vbs =
+          Map.singleton (module String) name @@ make_vb ~loc ~name ~name_expr:expr ~hint_label:None
+        in
         {
-          default_result with
+          vbs;
           typ = No_grad_tensor_intro { name; name_expr = expr };
           expr = A.Exp.ident ~loc:str_loc { txt = Lident name; loc = str_loc };
+          array_opt_of_code = None;
+          slot = Undet;
         }
     | { pexp_desc = Pexp_array _; _ }
     | { pexp_desc = Pexp_construct ({ txt = Lident "::"; _ }, _); _ } ->
@@ -852,9 +859,8 @@ let translate (expr : expression) : result =
             }
         | Code ->
             {
-              default_result with
+              res1 with
               typ = Array;
-              slot = res1.slot;
               expr =
                 Ast_builder.Default.pexp_extension ~loc
                 @@ Location.error_extensionf ~loc
@@ -878,7 +884,7 @@ let translate (expr : expression) : result =
                 @@ Location.error_extensionf ~loc
                      "ppx_ocannl %%cd: only tensor nodes (e.g. `.value` or `.grad`) can be merged";
             }
-        | Grad_of_tensor t -> { res1 with vbs = no_vbs; typ = Merge_grad t }
+        | Grad_of_tensor t -> { res1 with typ = Merge_grad t }
         | Merge_value _ | Merge_grad _ ->
             {
               res1 with
diff --git a/lib/ppx_shared.ml b/lib/ppx_shared.ml
@@ -241,7 +241,7 @@ let reduce_vbss = List.reduce_exn ~f:(Map.merge_skewed ~combine:(fun ~key:_ _v1
 let expr_expander_with_punning translate ~loc ~path:_ payload =
   match payload with
   | { pexp_desc = Pexp_let (recflag, bindings, body); _ } ->
-      (* We are at the %op annotation level: do not tranlsate the body. *)
+      (* We are at the %op/%cd annotation level: do not tranlsate the body. *)
       let vbss, bindings =
         List.unzip
         @@ List.map bindings ~f:(fun vb ->
diff --git a/lib/syntax_extensions.md b/lib/syntax_extensions.md
@@ -255,7 +255,7 @@ When an extension is over a wildcard (ignore result) binding: `let%cd _ = ...` a
 
 Both `%cd` and `%op` syntaxes support inline declarations of tensors. For `%op` these are differentiable, for `%cd` non-differentiable tensors. A declaration site uses the string syntax, the content of the string is the is bound to the newly created tensor, and the string itself functions equivalently to using the newly introduced identifier. The scope of the binding is the full scope of the extension point, even if the declaring string appeared in the body of a function that's inside the extension point scope (except for `%op` there is a special case of `~config` labeled argument discussed below). The first element of the label of the created tensor is the string that introduced it.
 
-For `%cd`, the declaration is (currently) only allowed on the left-hand-side, i.e. in the assigned-to position, of an assignment. If possible, one of the tensors on the right-hand-side is picked to provide additional label information. In particular, tensors that are function parameters inside the scope of the extension point, cannot be picked to provide label information, as they would escape their scope at the point the tensor is created. Example showing two tensor nodes declared inline, both of them include the label of the param `p` in their labels:
+For `%cd`, inline declarations are allowed both in the assigned-to position (left-hand side) of assignments and in standalone tensor expressions. When used in assignments, one of the tensors on the right-hand-side is picked to provide additional label information if possible. In particular, tensors that are function parameters inside the scope of the extension point, cannot be picked to provide label information, as they would escape their scope at the point the tensor is created. Inline declarations are still prohibited within the right-hand side of assignments to discourage over-use in locations with less label information. Example showing two tensor nodes declared inline, both of them include the label of the param `p` in their labels:
 
 ```ocaml
 let sgd_one ~learning_rate ?(momentum = 0.0) ?(weight_decay = 0.0) ?(nesterov = false) p =
@@ -267,6 +267,21 @@ let sgd_one ~learning_rate ?(momentum = 0.0) ?(weight_decay = 0.0) ?(nesterov =
     p =- learning_rate *. sgd_delta]
 ```
 
+Inline declarations can also be used outside of assignments for creating non-differentiable tensors, to mimic the behavior of `%op` but without the burden of initialization that a parameter would introduce:
+
+```ocaml
+  let%cd mlp_result = mlp "point" in
+  let result_routine =
+    Train.to_routine (module Backend) sgd_routine.context IDX.empty
+      [%cd ~~("mlp infer"; mlp_result.forward)]
+  in
+  let callback (x, y) =
+    Tn.set_values point [| x; y |];
+    Train.run result_routine;
+    Float.(mlp_result.@[0] >= 0.)
+  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:
 
 ```ocaml
diff --git a/test/einsum/moons_demo_variant.ml b/test/einsum/moons_demo_variant.ml
@@ -59,13 +59,13 @@ let () =
   let%op learning_rate = 0.1 *. ((2 *. !..steps) - !@step_n) /. !..steps in
   Train.set_hosted learning_rate.value;
   let sgd = Train.sgd_update ~learning_rate ~weight_decay scalar_loss in
-  let init = Train.to_routine (module Backend) ctx bindings init_params in
+  let init_routine = Train.to_routine (module Backend) ctx bindings init_params in
   let sgd_routine =
-    Train.to_routine (module Backend) init.context bindings (Asgns.sequence [ update; sgd ])
+    Train.to_routine (module Backend) init_routine.context bindings (Asgns.sequence [ update; sgd ])
   in
   let step_ref = IDX.find_exn sgd_routine.bindings step_n in
   step_ref := 0;
-  Train.run init;
+  Train.run init_routine;
   for _epoch = 1 to epochs do
     Train.sequential_loop sgd_routine.bindings ~f:(fun () ->
         Train.run sgd_routine;
