Embedding of dimensions in tensor expressions: %cd syntax extension and row variable test, by Claude Opus

lukstafi · lukstafi · commit 58a0af5f3c24 · 2025-09-02T14:36:20.000+02:00
Summary by Claude:
The implementation is now complete. We have:

  1. ✅ Implemented the apply_env_step function in shape.ml to update
  delayed_var_ref fields
  2. ✅ Added helper functions to Row module for extracting resolved
  dimensions
  3. ✅ Updated the delayed_var_ref vars with resolved dimensions
  during shape inference
  4. ✅ Implemented parsing of capture_dims list in ppx_op.ml for
  einsum operations
  5. ✅ Created Indexing.variable_ref objects and bound them in
  ppx_op.ml
  6. ✅ Added support for capture_dims in ppx_cd.ml for %cd syntax
  7. ✅ Added test case for row variable capture that correctly shows
  the product of dimensions

  The feature for embedding dimensions in tensor expressions is now
  fully implemented and tested!
diff --git a/lib/ppx_cd.ml b/lib/ppx_cd.ml
@@ -1005,11 +1005,44 @@ let translate ?ident_label (expr : expression) : result =
           expr = [%expr einsum [%e spec] [%e res1.expr] [%e res2.expr]];
           array_opt_of_code = None;
         }
+    | [%expr
+        [%e? expr1]
+        *+ [%e? { pexp_desc = Pexp_constant (Pconst_string (spec_str, _, _)); _ }]
+           ([%e? { pexp_desc = Pexp_constant (Pconst_string _); _ } as head] :: [%e? rest])
+           [%e? expr2]]
+      when String.contains spec_str '>' ->
+        let capture_vbs, capture_dims_expr = collect_capture_labels ~loc head rest in
+        let res1 = loop ~proj_in_scope expr1 in
+        let res2 = loop ~proj_in_scope expr2 in
+        let spec = substitute_identifiers_in_einsum_spec ~loc spec_str in
+        let slot = List.hd_exn @@ List.sort [ res1.slot; res2.slot ] ~compare:compare_slots in
+        {
+          vbs = reduce_vbss [ res1.vbs; res2.vbs; capture_vbs ];
+          typ = Tensor;
+          slot;
+          expr = [%expr einsum ~capture_dims:[%e capture_dims_expr] [%e spec] [%e res1.expr] [%e res2.expr]];
+          array_opt_of_code = None;
+        }
     | [%expr [%e? expr1] ++ [%e? { pexp_desc = Pexp_constant (Pconst_string (spec_str, _, _)); _ }]]
       when String.contains spec_str '>' ->
         let res1 = loop ~proj_in_scope expr1 in
         let spec = substitute_identifiers_in_einsum_spec ~loc spec_str in
         { res1 with typ = Tensor; expr = [%expr einsum1 [%e spec] [%e res1.expr]] }
+    | [%expr
+        [%e? expr1]
+        ++ [%e? { pexp_desc = Pexp_constant (Pconst_string (spec_str, _, _)); _ }]
+           ([%e? { pexp_desc = Pexp_constant (Pconst_string _); _ } as head] :: [%e? rest])]
+      when String.contains spec_str '>' ->
+        let capture_vbs, capture_dims_expr = collect_capture_labels ~loc head rest in
+        let res1 = loop ~proj_in_scope expr1 in
+        let spec = substitute_identifiers_in_einsum_spec ~loc spec_str in
+        {
+          vbs = reduce_vbss [ res1.vbs; capture_vbs ];
+          typ = Tensor;
+          slot = res1.slot;
+          expr = [%expr einsum1 ~capture_dims:[%e capture_dims_expr] [%e spec] [%e res1.expr]];
+          array_opt_of_code = None;
+        }
     | [%expr [%e? expr1].grad] -> (
         let res1 = loop ~proj_in_scope expr1 in
         match res1.typ with
diff --git a/test/operations/test_einsum_capture.expected b/test/operations/test_einsum_capture.expected
@@ -1,16 +1,15 @@
-Retrieving commandline, environment, or config file variable ocannl_log_level
-Found 0, in the config file
 Dimension a: 2
 Dimension b: 3
 Dimension c: 4
 Dimension i: 5
 Dimension j: 7
-HERE: test/operations/test_einsum_capture.ml:39:21
+Row variable r (product of dims): 12
+HERE: test/operations/test_einsum_capture.ml:51:21
 ┌───────────────────────────┐
-│[25]: +_dim_calc shape 0:1 │
-│┌┬──────┐                  │
-│││axis 0│                  │
-│├┼──────┤                  │
-│││ 9.00 │                  │
-│└┴──────┘                  │
-└───────────────────────────┘
+│[41]: +_dim_calc shape 0:1 │
+│┌┬─────────┐               │
+│││axis 0   │               │
+│├┼─────────┤               │
+│││ 2.10e+1 │               │
+│└┴─────────┘               │
+└───────────────────────────┘
diff --git a/test/operations/test_einsum_capture.ml b/test/operations/test_einsum_capture.ml
@@ -33,7 +33,19 @@ let () =
   Stdio.printf "Dimension j: %s\n"
     (match j.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
 
-  let%op dim_calc = dim a + dim j in
+  (* Test capturing row variables *)
+  let%op x3 = { x3 = uniform1 (); o = [ 2; 3; 4 ] } in
+  let%op y3 = { y3 = uniform1 (); o = [ 3; 4; 5 ] } in
+  let%op z3 = x3 *+ "a..r..;..r..b=>ab" [ "r" ] y3 in
+  
+  (* Trigger shape inference *)
+  let ctx = Train.forward_once (module Backend) ~ctx z3 in
+  
+  (* Check if row variable was captured *)
+  Stdio.printf "Row variable r (product of dims): %s\n"
+    (match r.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
+  
+  let%op dim_calc = dim a + dim j + dim r in
   let _ctx = Train.forward_once (module Backend) ~ctx dim_calc in
 
   Train.printf ~here:[%here] ~with_code:false ~with_grad:false dim_calc
