A new convenience operation offsets and fix to ndarray (it's not composable so doesn't belong to O)

lukstafi · lukstafi · commit 82ea9158b225 · 2025-09-04T11:37:13.000+02:00
diff --git a/lib/operation.ml b/lib/operation.ml
@@ -404,7 +404,7 @@ let einsum1 ?(label = []) ?(capture_dims = []) spec =
     ~op_asn ~grad_asn ~label:("=>" :: label)
 
 module NDO_before_einmax1 = struct
-  let (+) ?label t1 t2 = add ?label ~grad_spec:Prohibit_grad t1 t2 ()
+  let ( + ) ?label t1 t2 = add ?label ~grad_spec:Prohibit_grad t1 t2 ()
   let where ?label t1 t2 t3 = where ?label ~grad_spec:Prohibit_grad t1 t2 t3 ()
   let not ?label t = not ?label ~grad_spec:Prohibit_grad t ()
   let ( < ) ?label t1 t2 = lt ?label ~grad_spec:Prohibit_grad t1 t2 ()
@@ -437,6 +437,9 @@ let tropical ?(label = []) ?(capture_dims = []) spec =
     ~compose_op:(Shape.Einsum (spec, capture_dims))
     ~op_asn ~grad_asn ~label:("@^=>+" :: label)
 
+(** A fully-shape-inferred tensor that is initialized with the offset of each cell. *)
+let offsets = Tensor.term ~fetch_op:Range_over_offsets ?init_data:None
+
 (** [range] is a 1D tensor of shape [upto], spans [0] inclusive, [upto] exclusive. *)
 let range ?(label = []) ?(grad_spec = Tensor.Prohibit_grad) ?axis_label upto =
   let result =
@@ -599,6 +602,7 @@ struct
   let einsum1 = einsum1 ~grad_spec:Grad_spec.grad_spec
   let einmax1 = einmax1 ~grad_spec:Grad_spec.grad_spec
   let tropical = tropical ~grad_spec:Grad_spec.grad_spec
+  let offsets = offsets ~grad_spec:Grad_spec.grad_spec
   let range = range ~grad_spec:Grad_spec.grad_spec
   let range_of_shape = range_of_shape ~grad_spec:Grad_spec.grad_spec
   let stop_gradient = stop_gradient
@@ -692,10 +696,11 @@ struct
     let ( <> ) ?label t1 t2 = ne ?label t1 t2 ()
     let embed_self_id = embed_self_id
     let einsum ?label ?capture_dims spec t1 t2 = einsum ?label ?capture_dims spec t1 t2 ()
+    let outer_sum ?label ?capture_dims spec t1 t2 = outer_sum ?label ?capture_dims spec t1 t2 ()
     let einsum1 ?label ?capture_dims spec t1 = einsum1 ?label ?capture_dims spec t1 ()
     let einmax1 ?label ?capture_dims spec t1 = einmax1 ?label ?capture_dims spec t1 ()
     let tropical ?label ?capture_dims spec t1 t2 = tropical ?label ?capture_dims spec t1 t2 ()
-    let ndarray = ndarray
+    let offsets ?label () = offsets ?label ()
     let uniform ?label () = uniform () ?label ()
     let uniform_at ?label counter = uniform_at ?label counter ()
     let uniform1 ?label () = uniform1 () ?label ()
diff --git a/lib/ppx_cd.ml b/lib/ppx_cd.ml
@@ -914,7 +914,7 @@ let translate ?ident_label (expr : expression) : result =
             })
     | { pexp_desc = Pexp_array _; _ }
     | { pexp_desc = Pexp_construct ({ txt = Lident "::"; _ }, _); _ } ->
-        { default_result with expr = ndarray_op expr }
+        { default_result with expr = ndarray_op ~ndarray_fn:[%expr NTDSL.ndarray] expr }
     | { pexp_desc = Pexp_ident { txt = Lident ("v" | "lhs"); _ }; _ } ->
         { default_result with typ = Array; slot = LHS }
     | { pexp_desc = Pexp_ident { txt = Lident "g"; _ }; _ } ->
diff --git a/lib/ppx_op.ml b/lib/ppx_op.ml
@@ -45,7 +45,7 @@ let operators =
       ("embed_self_id", "embed_self_id");
       ("einsum", "einsum");
       ("einsum1", "einsum1");
-      ("ndarray", "ndarray");
+      ("offsets", "offsets");
       ("uniform", "uniform");
       ("uniform_at", "uniform_at");
       ("uniform1", "uniform1");
@@ -325,7 +325,7 @@ let rec translate ~num_configs ~is_toplevel ~opt_label ?label expr =
                  "ppx_ocannl %%op: record field label must be a simple identifier" ))
   | { pexp_desc = Pexp_array _; _ }
   | { pexp_desc = Pexp_construct ({ txt = Lident "::"; _ }, _); _ } ->
-      (no_vbs, ndarray_op ?label expr)
+      (no_vbs, ndarray_op ?label ~ndarray_fn:[%expr TDSL.ndarray] expr)
   | [%expr !.[%e? expr1]] ->
       (* Hardcoding the patterns for (!.), (!..), and ( **. ) to avoid treating the constants as
          already tensors. *)
diff --git a/lib/ppx_shared.ml b/lib/ppx_shared.ml
@@ -256,13 +256,17 @@ let assignment_ops =
 let einsum_binary_ops =
   Hashtbl.of_alist_exn
     (module String)
-    [ ("+*", fun loc -> [%expr einsum]); ("@^+", fun loc -> [%expr tropical]) ]
+    [
+      ("+*", fun loc -> [%expr einsum]);
+      ("@^+", fun loc -> [%expr tropical]);
+      ("+++", fun loc -> [%expr outer_sum]);
+    ]
 
 let einsum_unary_ops =
-      Hashtbl.of_alist_exn
-        (module String)
-        [ ("++", fun loc -> [%expr einsum1]); ("@^^", fun loc -> [%expr einmax1]) ]
-    
+  Hashtbl.of_alist_exn
+    (module String)
+    [ ("++", fun loc -> [%expr einsum1]); ("@^^", fun loc -> [%expr einmax1]) ]
+
 let is_primitive_op op_ident =
   List.exists ~f:(Fn.flip Hashtbl.mem op_ident) [ ternary_ops; unary_ops; binary_ops ]
 
@@ -313,11 +317,11 @@ let translate_str translate ({ pstr_desc; pstr_loc = loc; _ } as str) =
 let str_expander_with_punning translate ~loc ~path (payload : structure_item list) =
   flatten_str ~loc ~path @@ List.map payload ~f:(translate_str translate)
 
-let ndarray_op ?axis_labels ?label expr =
+let ndarray_op ?axis_labels ?label ~ndarray_fn expr =
   let loc = expr.pexp_loc in
   let values, batch_dims, output_dims, input_dims = ndarray_constant expr in
   let edims dims = Ast_builder.Default.elist ~loc dims in
-  let w_val = [%expr ndarray [%e values]] in
+  let w_val = [%expr [%e ndarray_fn] [%e values]] in
   let op =
     match (axis_labels, label) with
     | None, None -> w_val
diff --git a/test/ppx/test_ppx_op_expected.ml b/test/ppx/test_ppx_op_expected.ml
@@ -32,13 +32,13 @@ let y2 =
         ((( *. ) ?label:None) x1 hey3) x2
 let a =
   let open! TDSL.O in
-    ((ndarray
+    ((TDSL.ndarray
         [|(Float.of_int 1);(Float.of_int 2);(Float.of_int 3);(Float.of_int 4);(
           Float.of_int 5);(Float.of_int 6)|]) ~label:["a"]) ~batch_dims:[]
       ~input_dims:[3] ~output_dims:[2] ()
 let b =
   let open! TDSL.O in
-    ((ndarray
+    ((TDSL.ndarray
         [|(Float.of_int 7);(Float.of_int 8);(Float.of_int 9);(Float.of_int 10)|])
        ~label:["b"]) ~batch_dims:[2] ~input_dims:[] ~output_dims:[2] ()
 let y =
