New no-op shape inference spec Defined_by_cd_logic

lukstafi · lukstafi · commit 81a2d5e16f97 · 2025-11-21T11:00:20.000+01:00
diff --git a/tensor/operation.ml b/tensor/operation.ml
@@ -287,11 +287,12 @@ let interleave =
   let%cd op_asn ~t ~t1 ~t2 ~projections:_ =
     t =:+ id t1 ~logic:"... | ... -> ..., i => ... | ... -> ..., 2*i";
     t =+ id t2 ~logic:"... | ... -> ..., i => ... | ... -> ..., 2*i + 1"
-   in
+  in
   let%cd grad_asn ~t ~g:_ ~t1 ~t2 ~projections:_ =
-  t1.grad =+ id t.grad ~logic:"... | ... -> ..., 2*i => ... | ... -> ..., i";
-  t2.grad =+ id t.grad ~logic:"... | ... -> ..., 2*i + 1 => ... | ... -> ..., i"  in
-  Tensor.binop ~op_label:"<>" ~compose_op:Pointwise_bin ~op_asn ~grad_asn
+    t1.grad =+ id t.grad ~logic:"... | ... -> ..., 2*i => ... | ... -> ..., i";
+    t2.grad =+ id t.grad ~logic:"... | ... -> ..., 2*i + 1 => ... | ... -> ..., i"
+  in
+  Tensor.binop ~op_label:"interleave" ~compose_op:Defined_by_cd_logic ~op_asn ~grad_asn
 
 let threefry4x32_crypto =
   let module NTDSL = Initial_NTDSL in
diff --git a/tensor/shape.ml b/tensor/shape.ml
@@ -102,7 +102,11 @@ type delayed_var_ref = {
 let get_variable_ref ref_label =
   { var_ref = Ir.Indexing.{ ref_label; solved_dim = None }; var = `Not_set_yet }
 
-type compose_type = Pointwise_bin | Compose | Einsum of string * delayed_var_ref list
+type compose_type =
+  | Pointwise_bin
+  | Compose
+  | Einsum of string * delayed_var_ref list
+  | Defined_by_cd_logic
 [@@deriving sexp_of, equal]
 
 type transpose_type =
@@ -111,12 +115,14 @@ type transpose_type =
   | Permute of string * delayed_var_ref list
   | Batch_slice of Idx.static_symbol
   | Uint4x32_to_prec of Ir.Ops.prec Lazy.t
+  | Defined_by_cd_logic
 [@@deriving equal, sexp_of]
 
 type terminal_type = Data of Ir.Assignments.init_data | Fetch of Ir.Assignments.fetch_op
 [@@deriving equal, sexp_of]
 
-type ternary_type = Pointwise_tern | Compose_accumulate [@@deriving sexp, equal]
+type ternary_type = Pointwise_tern | Compose_accumulate | Defined_by_cd_logic
+[@@deriving sexp, equal]
 
 let identifier ~multichar =
   let open Angstrom in
@@ -277,6 +283,10 @@ let logic_to_spec = function
   | Transpose (Transpose, _) -> "T"
   | Transpose (Batch_slice _, _) -> "@|"
   | Transpose (Uint4x32_to_prec _, _) -> "U4x32"
+  | Broadcast (Defined_by_cd_logic, _, _)
+  | Transpose (Defined_by_cd_logic, _)
+  | Broadcast_tern (Defined_by_cd_logic, _, _, _) ->
+      "<cd_logic>"
   | Terminal _ -> "<terminal>"
 
 module Update_id = struct
@@ -371,8 +381,7 @@ let axes_spec_to_dims_bio ~sh_id ~row_var_env ~dim_var_env:_ ~f labels =
   let output = to_row `Output labels.bcast_output o_dims beg_o_dims in
   (batch, input, output)
 
-let einsum_slot_spec_to_dims_bio ~original_spec ~sh_id ~row_var_env ~dim_var_env labels
-    =
+let einsum_slot_spec_to_dims_bio ~original_spec ~sh_id ~row_var_env ~dim_var_env labels =
   let proj_env_update = ref @@ Row.dim_map_empty in
   let extras = ref [] in
   let f kind = function
@@ -840,6 +849,10 @@ let%debug4_sexp get_inequalities ({ shape = cur_sh; logic; id = _ } as _upd : up
               subr = sh3.output;
             };
         ] )
+  | Broadcast (Defined_by_cd_logic, _, _)
+  | Transpose (Defined_by_cd_logic, _)
+  | Broadcast_tern (Defined_by_cd_logic, _, _, _) ->
+      (Row.dim_map_empty, [])
   | Transpose (Batch_slice { static_range; static_symbol }, sh) ->
       Hash_set.remove unused_shapes sh.id;
       let slice_v = get_var () in
@@ -915,12 +928,12 @@ let%debug4_sexp get_inequalities ({ shape = cur_sh; logic; id = _ } as _upd : up
       let dim_var_env = Hashtbl.create (module String) in
 
       let extras_rhs, proj_env_rhs, (b_rhs, i_rhs, o_rhs) =
-        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:sh.id ~row_var_env
-          ~dim_var_env ls_rhs
+        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:sh.id ~row_var_env ~dim_var_env
+          ls_rhs
       in
       let extras_lhs, proj_env_lhs, (b_lhs, i_lhs, o_lhs) =
-        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:cur_sh.id ~row_var_env
-          ~dim_var_env ls_lhs
+        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:cur_sh.id ~row_var_env ~dim_var_env
+          ls_lhs
       in
       (* Bind delayed_var_refs to the variables after they are created *)
       let extras_dim_refs =
@@ -1134,16 +1147,16 @@ let%debug4_sexp get_inequalities ({ shape = cur_sh; logic; id = _ } as _upd : up
       let row_var_env = Hashtbl.create (module String) in
       let dim_var_env = Hashtbl.create (module String) in
       let extras_rhs1, proj_env_rhs1, (b_rhs1, i_rhs1, o_rhs1) =
-        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:sh1.id ~row_var_env
-          ~dim_var_env ls_rhs1
+        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:sh1.id ~row_var_env ~dim_var_env
+          ls_rhs1
       in
       let extras_rhs2, proj_env_rhs2, (b_rhs2, i_rhs2, o_rhs2) =
-        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:sh2.id ~row_var_env
-          ~dim_var_env ls_rhs2
+        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:sh2.id ~row_var_env ~dim_var_env
+          ls_rhs2
       in
       let extras_lhs, proj_env_lhs, (b_lhs, i_lhs, o_lhs) =
-        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:cur_sh.id ~row_var_env
-          ~dim_var_env ls_lhs
+        einsum_slot_spec_to_dims_bio ~original_spec:spec ~sh_id:cur_sh.id ~row_var_env ~dim_var_env
+          ls_lhs
       in
       (* Bind delayed_var_refs to the variables after they are created *)
       (* TODO: refactor to avoid duplication with the one for unary einsum *)
@@ -1735,6 +1748,10 @@ let fresh_proj_ids update =
   in
   fresh_shape update.shape;
   (match update.logic with
+  | Broadcast (Defined_by_cd_logic, _, _)
+  | Transpose (Defined_by_cd_logic, _)
+  | Broadcast_tern (Defined_by_cd_logic, _, _, _) ->
+      ()
   | Terminal _ -> ()
   | Transpose (_, sh) -> fresh_shape sh
   | Broadcast (_, sh1, sh2) ->
@@ -1777,6 +1794,12 @@ let%debug4_sexp derive_projections (update_step : update_step) : Idx.projections
   let lhs = update_step.shape in
   let rhs =
     match update_step.logic with
+    | Broadcast (Defined_by_cd_logic, _, _)
+    | Transpose (Defined_by_cd_logic, _)
+    | Broadcast_tern (Defined_by_cd_logic, _, _, _) ->
+        raise
+        @@ Utils.User_error
+             "Defined_by_cd_logic: use explicit ~logic annotations when defining this operation"
     | Terminal _ -> []
     | Transpose (_, sh) -> [ sh ]
     | Broadcast (_, sh1, sh2) -> [ sh1; sh2 ]
@@ -1830,7 +1853,7 @@ let make ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_ax
     match kind with
     | `Batch | `Input -> get_dim ~d ()
     | `Output ->
-        if not known_no_batch && num_dim1_output = 1 && d = 1 then
+        if (not known_no_batch) && num_dim1_output = 1 && d = 1 then
           let label = debug_name ^ "_output" in
           get_dim ~d ~label ()
         else get_dim ~d ()
diff --git a/tensor/shape.mli b/tensor/shape.mli
@@ -113,20 +113,22 @@ type compose_type =
           multiply). *)
   | Einsum of string * delayed_var_ref list
       (** The binary "einsum" syntax: RHS1;RHS2=>LHS, where RHSi, LHS are labels specifications.
-          OCANNL's extended einsum notation supports both axis variables and row variables.
-          The [axis_labels] use pseudo-labels local to the notation, to line up the axes and row
+          OCANNL's extended einsum notation supports both axis variables and row variables. The
+          [axis_labels] use pseudo-labels local to the notation, to line up the axes and row
           variables. The symmetric difference / disjunctive union of RHS1 and RHS2's pseudo-labels
           should be equal to LHS pseudo-labels.
 
           Unlike [Pointwise_bin] and [Compose], einsum operations use equations only (not
-          inequalities), so they do NOT permit broadcasting. This makes einsum more restrictive
-          but also more precise for operations where exact shape matching is required.
+          inequalities), so they do NOT permit broadcasting. This makes einsum more restrictive but
+          also more precise for operations where exact shape matching is required.
 
           The optional {!Ir.Indexing.variable_ref}s will capture the solutions of the dimensions
           corresponding to the specification labels equal to [ref_label] of a reference.
 
           Note: The "right-hand-side" is on the left! I.e. the syntax is "rhs=>lhs",
           "rhs1;rhs2=>lhs". *)
+  | Defined_by_cd_logic
+      (** A placeholder for operations where the shape logic is defined by the %cd extension. *)
 [@@deriving sexp_of, equal]
 
 type transpose_type =
@@ -136,8 +138,8 @@ type transpose_type =
       (** The unary "einsum" syntax: RHS1=>LHS.
 
           Unlike [Pointwise_un], permute operations use equations only (not inequalities), so they
-          do NOT permit broadcasting. This makes permute more restrictive but also more precise
-          for operations where exact shape matching is required.
+          do NOT permit broadcasting. This makes permute more restrictive but also more precise for
+          operations where exact shape matching is required.
 
           The optional {!Ir.Indexing.variable_ref}s will capture the solutions of the dimensions
           corresponding to the specification labels equal to [ref_label] of a reference. *)
@@ -146,13 +148,17 @@ type transpose_type =
       (** Converts precision in a bit-effient way, with a corresponding conversion in total number
           of elements. Currently, assumes the incoming tensor (RHS) has just a single axis to not
           force unnecessary minimum sizes on output axes. *)
+  | Defined_by_cd_logic
+      (** A placeholder for operations where the shape logic is defined by the %cd extension. *)
 [@@deriving equal, sexp_of]
 
 (** If you miss expressivity here, leave a note on
     {{:https://github.com/ahrefs/ocannl/issues/305}issue 305}. *)
 type ternary_type =
   | Pointwise_tern  (** As in the operation [Where]. *)
   | Compose_accumulate  (** As in the operation [FMA]. *)
+  | Defined_by_cd_logic
+      (** A placeholder for operations where the shape logic is defined by the %cd extension. *)
 [@@deriving equal, sexp_of]
 
 (** Extracts any available shape information from the initialization or fetch. *)
@@ -203,8 +209,8 @@ type logic =
           [s1], hence the name. *)
   | Broadcast_tern of ternary_type * t * t * t  (** Matches the shapes for a ternary operation. *)
   | Terminal of { is_param : bool; logic : terminal_type }
-      (** Extracts any available shape information from the initialization. 
-          The [is_param] field indicates if this is a parameter tensor that requires gradients. *)
+      (** Extracts any available shape information from the initialization. The [is_param] field
+          indicates if this is a parameter tensor that requires gradients. *)
 [@@deriving equal, sexp_of]
 
 type update_id [@@deriving equal, compare, hash, sexp]
diff --git a/test/operations/dune b/test/operations/dune
@@ -389,3 +389,14 @@
  (preprocess
   (pps ppx_here ppx_expect ppx_inline_test ppx_ocannl))
  (modes best))
+
+(test
+ (name test_interleave)
+ (package neural_nets_lib)
+ (deps
+  ocannl_config
+  (env_var OCANNL_BACKEND))
+ (modules test_interleave)
+ (libraries base ocannl stdio)
+ (preprocess
+  (pps ppx_here ppx_ocannl ppx_sexp_conv)))
diff --git a/test/operations/test_interleave.ml b/test/operations/test_interleave.ml
@@ -0,0 +1,25 @@
+open Base
+open Ocannl
+open Nn_blocks.DSL_modules
+
+let () =
+  Tensor.unsafe_reinitialize ();
+  let%op t1 = [ 1.0; 2.0; 3.0 ] in
+  let%op t2 = [ 4.0; 5.0; 6.0 ] in
+  let t3 = Operation.interleave t1 t2 () in
+
+  (* t3 should be [1.0; 4.0; 2.0; 5.0; 3.0; 6.0] *)
+  let ctx = Context.auto () in
+
+  try
+    let _ctx = Train.forward_once ctx t3 in
+    Stdio.printf "Test failed! Expected error was not raised.\n";
+    Stdlib.exit 1
+  with Utils.User_error msg ->
+    if
+      String.equal msg
+        "Defined_by_cd_logic: use explicit ~logic annotations when defining this operation"
+    then Stdio.printf "Test passed! Caught expected error: %s\n" msg
+    else (
+      Stdio.printf "Test failed! Caught unexpected error: %s\n" msg;
+      Stdlib.exit 1)
