In progress: New shape specification functionality: add equality constraints via captured variables

lukstafi · lukstafi · commit 9dcfcaad39e4 · 2025-09-03T11:06:14.000+02:00
Move to using `Shape`-level `delayed_var_ref` for variable capture. Extend Shape API with `set_dim` and `set_equal` introducing equality constraints, where "equal" between a row and a dimension is interpreted via `Total_elems` -- not by assuming the row is that single axis!
diff --git a/lib/operation.ml b/lib/operation.ml
@@ -80,8 +80,6 @@ let matmul ?(label = []) =
   let%cd op_asn ~v ~t1 ~t2 ~projections = v =:+ v1 * v2 in
   mul Compose ~op_asn ~label:("*" :: label)
 
-let capture_dims_to_refs = List.map ~f:(fun var_ref -> { Shape.var_ref; var = `Not_set_yet })
-
 (** Similar to the explicit mode of [numpy.einsum], the binary variant. Can compute various forms of
     matrix multiplication, inner and outer products, etc.
 
@@ -94,9 +92,7 @@ let einsum ?(label = []) ?(capture_dims = []) spec =
     g1 =+ g * v2;
     g2 =+ v1 * g
   in
-  Tensor.binop ~label:(";=>" :: label)
-    ~compose_op:(Einsum (spec, capture_dims_to_refs capture_dims))
-    ~op_asn ~grad_asn
+  Tensor.binop ~label:(";=>" :: label) ~compose_op:(Einsum (spec, capture_dims)) ~op_asn ~grad_asn
 
 (** Like [einsum], but adds instead than multiplying the resulting values. *)
 let outer_sum ?(label = []) ?(capture_dims = []) spec =
@@ -106,9 +102,7 @@ let outer_sum ?(label = []) ?(capture_dims = []) spec =
     g1 =+ g;
     g2 =+ g
   in
-  Tensor.binop ~label:(";=>+" :: label)
-    ~compose_op:(Einsum (spec, capture_dims_to_refs capture_dims))
-    ~op_asn ~grad_asn
+  Tensor.binop ~label:(";=>+" :: label) ~compose_op:(Einsum (spec, capture_dims)) ~op_asn ~grad_asn
 
 (** Similar to the explicit mode of [numpy.einsum], the unary variant. Can permute axes, extract
     diagonals, compute traces etc.
@@ -120,7 +114,7 @@ let einsum1 ?(label = []) ?(capture_dims = []) spec =
   let%cd op_asn ~v ~t1 ~projections = v =:+ v1 in
   let%cd grad_asn ~t:_ ~g ~t1 ~projections = g1 =+ g in
   Tensor.unop
-    ~transpose_op:(Shape.Permute (spec, capture_dims_to_refs capture_dims))
+    ~transpose_op:(Shape.Permute (spec, capture_dims))
     ~op_asn ~grad_asn ~label:("=>" :: label)
 
 module NDO_before_pow = struct
@@ -471,8 +465,8 @@ let embed_self_id ?grad_spec ?(label = []) () =
     ~input_dims:[] ~output_dims:[ 1 ] ()
 
 let embed_dim ?grad_spec ?(label = []) variable_ref =
-  Tensor.term ~fetch_op:(Embed_dim variable_ref) ?grad_spec ~label:("!@self_id" :: label)
-    ~batch_dims:[] ~input_dims:[] ~output_dims:[ 1 ] ()
+  Tensor.term ~fetch_op:(Embed_dim variable_ref.Shape.var_ref) ?grad_spec
+    ~label:("!@self_id" :: label) ~batch_dims:[] ~input_dims:[] ~output_dims:[ 1 ] ()
 
 let uniform ?grad_spec () =
   uint4x32_to_prec_uniform ?grad_spec
@@ -674,6 +668,7 @@ end
 module DSL_modules = struct
   module Shape = Shape
   module Tensor = Tensor
+
   module TDSL = Make_DSL (struct
     let grad_spec = Tensor.If_needed
   end)
diff --git a/lib/ppx_shared.ml b/lib/ppx_shared.ml
@@ -334,13 +334,7 @@ let collect_capture_labels ~loc head rest =
   in
   let capture_refs, capture_bindings =
     List.map capture_labels ~f:(fun (loc, label) ->
-        let ref_expr =
-          [%expr
-            {
-              Ir.Indexing.ref_label = [%e Ast_builder.Default.estring ~loc label];
-              solved_dim = None;
-            }]
-        in
+        let ref_expr = [%expr Shape.get_variable_ref [%e Ast_builder.Default.estring ~loc label]] in
         let binding =
           Ast_builder.Default.value_binding ~loc
             ~pat:(Ast_builder.Default.pvar ~loc label)
diff --git a/lib/row.ml b/lib/row.ml
@@ -169,6 +169,9 @@ let phantom_row_id = row_id ~sh_id:(-1) ~kind:`Output
 type t = { dims : dim list; bcast : bcast; id : row_id } [@@deriving equal, hash, compare, sexp]
 type row = t [@@deriving equal, sexp]
 
+let get_row_for_var ?(row_id = phantom_row_id) v =
+  { dims = []; bcast = Row_var { v; beg_dims = [] }; id = row_id }
+
 let dims_label_assoc dims =
   let f = function Var { label = Some l; _ } as d -> Some (l, d) | _ -> None in
   List.filter_map dims.dims ~f
diff --git a/lib/row.mli b/lib/row.mli
@@ -61,6 +61,7 @@ type bcast =
 type t = { dims : dim list; bcast : bcast; id : row_id } [@@deriving equal, hash, compare, sexp]
 
 val dims_label_assoc : t -> (string * dim) list
+val get_row_for_var : ?row_id:row_id -> row_var -> t
 
 type environment [@@deriving sexp_of]
 type error_trace = ..
diff --git a/lib/shape.ml b/lib/shape.ml
@@ -98,6 +98,9 @@ type delayed_var_ref = {
 }
 [@@deriving equal, sexp_of]
 
+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
 [@@deriving sexp_of, equal]
 
@@ -685,6 +688,46 @@ let state = ref Row.empty_env
 let active_update_steps = ref []
 let active_constraints = ref []
 
+let set_dim delayed_var_ref dim =
+  match delayed_var_ref with
+  | { var_ref = { solved_dim = Some dim2; _ }; _ } when dim2 = dim -> ()
+  | { var_ref = { solved_dim = Some dim2; ref_label; _ }; _ } ->
+      raise
+      @@ Row.Shape_error
+           ( "Cannot set dimension for variable reference with label " ^ ref_label,
+             [ Row.Dim_mismatch [ Row.get_dim ~d:dim2 (); Row.get_dim ~d:dim () ] ] )
+  | { var_ref = { solved_dim = None; _ }; var = `Not_set_yet } ->
+      delayed_var_ref.var_ref.solved_dim <- Some dim
+  | { var_ref = { solved_dim = None; _ }; var = `Dim dim_var } ->
+      delayed_var_ref.var_ref.solved_dim <- Some dim;
+      active_constraints :=
+        Row.Dim_eq { d1 = Row.Var dim_var; d2 = Row.get_dim ~d:dim () } :: !active_constraints
+  | { var_ref = { solved_dim = None; _ }; var = `Row row_var } ->
+      delayed_var_ref.var_ref.solved_dim <- Some dim;
+      active_constraints :=
+        Row.Rows_constr
+          {
+            (* TODO: actually, the Row.row_id should be the one of the shape that the row variable
+               is in, should be stored in `Row and in env_row_var. *)
+            r = [ Row.get_row_for_var row_var ];
+            constr = Total_elems { numerator = Num_elems dim; divided_by = [] };
+          }
+        :: !active_constraints
+
+let set_equal delayed_ref1 delayed_ref2 =
+  match delayed_ref1, delayed_ref2 with
+  | { var_ref = { solved_dim = Some dim1; _ }; _ },
+    { var_ref = { solved_dim = Some dim2; _ }; _ } ->
+      if dim1 = dim2 then ()
+      else
+        raise
+        @@ Row.Shape_error
+             ( "Cannot set equal dimensions for variable references with different values",
+               [ Row.Dim_mismatch [ Row.get_dim ~d:dim1 (); Row.get_dim ~d:dim2 () ] ] )
+  | _ -> 
+    (* FIXME: NOT IMPLEMENTED YET *)
+    ()
+
 let unsafe_reinitialize () =
   update_uid := 0;
   state := Row.empty_env;
diff --git a/lib/shape.mli b/lib/shape.mli
@@ -86,6 +86,24 @@ type delayed_var_ref = {
 }
 [@@deriving equal, sexp_of]
 
+val get_variable_ref : string -> delayed_var_ref
+(** Returns a fully unset variable reference with the given label. *)
+
+val set_dim : delayed_var_ref -> int -> unit
+(** Sets the dimension resp. total elements of the dim resp. row variable reference to the given
+    value. This will propagate through shape inference.
+
+    For row variables, this means the product of the dimensions, via the [Total_elems] constraint.
+*)
+
+val set_equal : delayed_var_ref -> delayed_var_ref -> unit
+(** Sets the two variable references to be equal (in some sense). This will propagate through shape
+    inference.
+
+    When both references are dimension variables or both are row variables, this means they are
+    precisely equal. When one is a dimension variable and the other is a row variable, this means
+    they have the same number of total elements. *)
+
 type compose_type =
   | Pointwise_bin
       (** NumPy-style broadcast matching batch, input and output axes, e.g. as in [s1 + s2]. *)
diff --git a/test/operations/test_einsum_capture.ml b/test/operations/test_einsum_capture.ml
@@ -13,11 +13,11 @@ let () =
 
   (* Check if dimensions were captured *)
   Stdio.printf "Dimension a: %s\n"
-    (match a.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
+    (match a.var_ref.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
   Stdio.printf "Dimension b: %s\n"
-    (match b.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
+    (match b.var_ref.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
   Stdio.printf "Dimension c: %s\n"
-    (match c.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
+    (match c.var_ref.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
 
   let%op x2 = { x2 = uniform1 (); o = [ 5; 7 ] } in
   (* Manually call einsum1 with capture_dims for now *)
@@ -28,9 +28,9 @@ let () =
 
   (* Check if dimensions were captured *)
   Stdio.printf "Dimension i: %s\n"
-    (match i.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
+    (match i.var_ref.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
   Stdio.printf "Dimension j: %s\n"
-    (match j.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
+    (match j.var_ref.solved_dim with Some d -> Int.to_string d | None -> "not resolved");
 
   (* Test capturing row variables *)
   let%op x3 = { x3 = uniform1 (); o = [ 2; 3; 4 ] } in
@@ -42,7 +42,7 @@ let () =
 
   (* 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");
+    (match r.var_ref.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
