Embedding of dimensions in tensor expressions: track variables so references can be updated by an upcoming apply_env_step

lukstafi · lukstafi · commit f650c0ce6223 · 2025-09-02T10:32:24.000+02:00
diff --git a/lib/operation.ml b/lib/operation.ml
@@ -80,6 +80,8 @@ 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.
 
@@ -92,7 +94,9 @@ let einsum ?(label = []) ?(capture_dims = []) spec =
     g1 =+ g * v2;
     g2 =+ v1 * g
   in
-  Tensor.binop ~label:(";=>" :: label) ~compose_op:(Einsum (spec, capture_dims)) ~op_asn ~grad_asn
+  Tensor.binop ~label:(";=>" :: label)
+    ~compose_op:(Einsum (spec, capture_dims_to_refs capture_dims))
+    ~op_asn ~grad_asn
 
 (** Like [einsum], but adds instead than multiplying the resulting values. *)
 let outer_sum ?(label = []) ?(capture_dims = []) spec =
@@ -102,7 +106,9 @@ let outer_sum ?(label = []) ?(capture_dims = []) spec =
     g1 =+ g;
     g2 =+ g
   in
-  Tensor.binop ~label:(";=>+" :: label) ~compose_op:(Einsum (spec, capture_dims)) ~op_asn ~grad_asn
+  Tensor.binop ~label:(";=>+" :: label)
+    ~compose_op:(Einsum (spec, capture_dims_to_refs 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.
@@ -114,7 +120,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))
+    ~transpose_op:(Shape.Permute (spec, capture_dims_to_refs capture_dims))
     ~op_asn ~grad_asn ~label:("=>" :: label)
 
 module NDO_before_pow = struct
diff --git a/lib/shape.ml b/lib/shape.ml
@@ -92,13 +92,19 @@ let row_of_kind = function `Batch -> batch | `Input -> input | `Output -> output
 type deduce_within_shape = Not_constrained | Input_equals_output
 [@@deriving compare, sexp, variants]
 
-type compose_type = Pointwise_bin | Compose | Einsum of string * Idx.variable_ref list
+type delayed_var_ref = {
+  var_ref : Ir.Indexing.variable_ref;
+  mutable var : [ `Row of Row.row_var | `Dim of Row.dim_var | `Not_set_yet ];
+}
+[@@deriving equal, sexp_of]
+
+type compose_type = Pointwise_bin | Compose | Einsum of string * delayed_var_ref list
 [@@deriving sexp_of, equal]
 
 type transpose_type =
   | Transpose
   | Pointwise_un
-  | Permute of string * Idx.variable_ref list
+  | Permute of string * delayed_var_ref list
   | Batch_slice of Idx.static_symbol
   | Uint4x32_to_prec of Ir.Ops.prec Lazy.t
 [@@deriving equal, sexp_of]
@@ -705,6 +711,7 @@ let%debug4_sexp propagate_shapes (update_step : update_step) : unit =
   active_constraints := ineqs @ !active_constraints;
   let ineqs', env = Row.solve_inequalities ~stage:Row.Stage1 ineqs !state in
   let _debug_remaining_constraints : Row.constraint_ list = ineqs' in
+  (* FIXME: call apply_env_step instead *)
   iter_shapes update_step ~f:(apply_env_t env);
   state := env
 
@@ -725,6 +732,7 @@ let%debug4_sexp finish_inference (() : unit) : unit =
   let unsolved, env = Row.solve_inequalities ~stage:Stage7 unsolved env in
   assert (List.is_empty unsolved);
   let _active_update_steps : update_step list = !active_update_steps in
+  (* FIXME: call apply_env_step instead *)
   List.iter ~f:(iter_shapes ~f:(apply_env_t env)) !active_update_steps;
   let _applied_update_steps : update_step list = !active_update_steps in
   active_constraints := [];
diff --git a/lib/shape.mli b/lib/shape.mli
@@ -80,14 +80,20 @@ type t = {
 
 type deduce_within_shape = Not_constrained | Input_equals_output [@@deriving compare, sexp]
 
+type delayed_var_ref = {
+  var_ref : Ir.Indexing.variable_ref;
+  mutable var : [ `Row of Row.row_var | `Dim of Row.dim_var | `Not_set_yet ];
+}
+[@@deriving equal, sexp_of]
+
 type compose_type =
   | Pointwise_bin
       (** NumPy-style broadcast matching batch, input and output axes, e.g. as in [s1 + s2]. *)
   | Compose
       (** Compose the outputs of the second shape with the inputs of the first shape, i.e. the shape
           of [fun x -> s1(s2(x))], or [s1 * s2] where [*] is the inner product (e.g. matrix
           multiply). *)
-  | Einsum of string * Ir.Indexing.variable_ref list
+  | Einsum of string * delayed_var_ref list
       (** The binary "einsum" syntax: RHS1;RHS2=>LHS, where RHSi, LHS are labels specifications.
           Since OCANNL's extended einsum notation supports both axis variables and row variables, it
           makes other compose types redundant. The [axis_labels] use pseudo-labels local to the
@@ -104,7 +110,7 @@ type compose_type =
 type transpose_type =
   | Transpose  (** Swaps inputs and outputs of a shape, preserves batch axes. *)
   | Pointwise_un  (** Preserves the shape. *)
-  | Permute of string * Ir.Indexing.variable_ref list
+  | Permute of string * delayed_var_ref list
       (** The unary "einsum" syntax: RHS1=>LHS.
 
           The optional {!Ir.Indexing.variable_ref}s will capture the solutions of the dimensions
