Untested: allow non-literal specification strings for einsum-like operators (%cd and %op)

lukstafi · lukstafi · commit 48aecf5effd9 · 2025-09-04T15:08:03.000+02:00
diff --git a/lib/ppx_cd.ml b/lib/ppx_cd.ml
@@ -990,6 +990,24 @@ let translate ?ident_label (expr : expression) : result =
     | [%expr [%e? expr1] **. [%e? expr2]] ->
         let res1 = loop ~proj_in_scope expr1 in
         { res1 with typ = Tensor; expr = [%expr NTDSL.O.( **. ) [%e res1.expr] [%e expr2]] }
+    | [%expr
+        [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
+          [%e? expr1]
+          ([%e? { pexp_desc = Pexp_ident _; _ } as spec] [%e? expr2])]
+      when Hashtbl.mem einsum_binary_ops op_ident ->
+        let res1 = loop ~proj_in_scope expr1 in
+        let res2 = loop ~proj_in_scope expr2 in
+        let slot = List.hd_exn @@ List.sort [ res1.slot; res2.slot ] ~compare:compare_slots in
+        {
+          vbs = reduce_vbss [ res1.vbs; res2.vbs ];
+          typ = Tensor;
+          slot;
+          expr =
+            [%expr
+              [%e Hashtbl.find_exn einsum_binary_ops op_ident loc]
+                [%e spec] [%e res1.expr] [%e res2.expr]];
+          array_opt_of_code = None;
+        }
     | [%expr
         [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
           [%e? expr1]
@@ -1031,7 +1049,19 @@ let translate ?ident_label (expr : expression) : result =
                 ~capture_dims:[%e capture_dims_expr] [%e spec] [%e res1.expr] [%e res2.expr]];
           array_opt_of_code = None;
         }
-    | [%expr
+        | [%expr
+        [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
+          [%e? expr1]
+          [%e? { pexp_desc = Pexp_ident _; _ } as spec]]
+      when Hashtbl.mem einsum_unary_ops op_ident ->
+        let res1 = loop ~proj_in_scope expr1 in
+        {
+          res1 with
+          typ = Tensor;
+          expr =
+            [%expr [%e Hashtbl.find_exn einsum_unary_ops op_ident loc] [%e spec] [%e res1.expr]];
+        }
+        | [%expr
         [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
           [%e? expr1]
           [%e? { pexp_desc = Pexp_constant (Pconst_string (spec_str, _, _)); _ }]]
diff --git a/lib/ppx_op.ml b/lib/ppx_op.ml
@@ -183,6 +183,17 @@ let rec translate ~num_configs ~is_toplevel ~opt_label ?label expr =
         [%expr
           TDSL.number ?label:[%e opt_expr ~loc label] ~axis_label:[%e axis] (Float.of_int [%e i])]
       )
+  | [%expr
+      [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
+        [%e? expr1]
+        ([%e? { pexp_desc = Pexp_ident _; _ } as spec] [%e? expr2])]
+    when Hashtbl.mem einsum_binary_ops op_ident ->
+      let vbs1, e1 = loop expr1 in
+      let vbs2, e2 = loop expr2 in
+      ( reduce_vbss [ vbs1; vbs2 ],
+        [%expr
+          [%e Hashtbl.find_exn einsum_binary_ops op_ident loc]
+            ?label:[%e opt_expr ~loc label] [%e spec] [%e e1] [%e e2]] )
   | [%expr
       [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
         [%e? expr1]
@@ -195,6 +206,16 @@ let rec translate ~num_configs ~is_toplevel ~opt_label ?label expr =
         [%expr
           [%e Hashtbl.find_exn einsum_binary_ops op_ident loc]
             ?label:[%e opt_expr ~loc label] [%e spec] [%e e1] [%e e2]] )
+  | [%expr
+      [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
+        [%e? expr1]
+        [%e? { pexp_desc = Pexp_ident _; _ } as spec]]
+    when Hashtbl.mem einsum_unary_ops op_ident ->
+      let vbs1, e1 = loop expr1 in
+      ( vbs1,
+        [%expr
+          [%e Hashtbl.find_exn einsum_unary_ops op_ident loc]
+            ?label:[%e opt_expr ~loc label] [%e spec] [%e e1]] )
   | [%expr
       [%e? { pexp_desc = Pexp_ident { txt = Lident op_ident; _ }; _ }]
         [%e? expr1]
