Fix: Shape.set_terminal for parameters

lukstafi · lukstafi · commit 4e85a0f8c155 · 2025-09-11T13:13:30.000+02:00
diff --git a/docs/shape_inference.md b/docs/shape_inference.md
@@ -174,7 +174,7 @@ Simplification of an inequality, and constraint propagation, can generate more c
 
 * Stage 1 is online as tensors are composed, and conservatively performs unification and constraint propagation. Stages 2, 3, 4 are only performed once necessary: when projections or dimensions are requested.
 * Stage 2, when solving the constraints, substitutes dim variables in terminal shapes that do not have a LUB or other constraints, by dimension-1. (This is generalized at stage 6 to all variables.) (FIXME: reconsider this, see the algo for row variables: a new LUB can still be inferred.) Forces coefficients coming from precision byte sizes.
-* Stage 3, when solving the constraints, sets yet-unknown dimension and row variables in terminal shapes to their least upper bounds (if any), but for rows only if they don't have a `Total_elems 1` constraint. It substitutes row variables in terminal shapes that do not have a LUB by one axis if that's required to satisfy the variable's constraint. In Total_elems constraints with multiple row variables, it substitutes row variables originating from axes of non-output kind, and which do not have a LUB, by no-further-axes -- otherwise these constraints can be too hard to unlock.
+* Stage 3, when solving the constraints, sets yet-unknown dimension and row variables in terminal shapes to their least upper bounds LUB (if any), but for rows only if they don't have a `Total_elems 1` constraint. It substitutes row variables in terminal shapes that do not have a LUB by one axis if that's required to satisfy the variable's constraint. In Total_elems constraints with multiple row variables, it substitutes row variables originating from axes of non-output kind, and which do not have a LUB, by no-further-axes -- otherwise these constraints can be too hard to unlock.
 * Stage 4 sets yet-unknown dimensions with >1 lower bounds from direct accesses, to their LUBs if they have any. It substitutes row variables in terminal shapes that do not have a LUB by no-further-axes. (This is generalized at stage 6 to all variables.) At this stage, we inject `Shape_row` constraints into the inequalities, so that we can re-process the variables of interest without traversing the whole environment.
 * Stage 5 addresses `Total_elems` and `Exact` constraints with yet-unknown row variables. For `Total_elems` and a single row variable: if the constraint can be satisfied by assuming the row variable is no-further-axes, it sets the row variable to `Broadcastable`, otherwise it sets it to one axis of the required dimension. For multiple row variables, if one is of the Output kind, sets the other variables to no-further-axes, and retries.
 * Stage 6 sets row variables in the remaining inequalities and updated shapes to no-further-axes values. This can unlock further between-axis inequalities because of row variables sandwiched between leftmost axes from their side of the inequality and rightmost axes from the other side of the inequality. In row constraints, this also unlocks inference for the embedded dim variables.
diff --git a/tensor/shape.ml b/tensor/shape.ml
@@ -1377,6 +1377,23 @@ let set_equal delayed_ref1 delayed_ref2 =
           }
         :: !active_constraints
 
+let set_terminal sh =
+  let get_origin kind =
+    Row.
+      {
+        lhs_name = sh.debug_name;
+        lhs_kind = kind;
+        rhs_name = "(parameter)";
+        rhs_kind = kind;
+        operation = Some "set_terminal";
+      }
+  in
+  active_constraints :=
+    Row.Terminal_row (sh.batch, [ get_origin `Batch ])
+    :: Row.Terminal_row (sh.input, [ get_origin `Input ])
+    :: Row.Terminal_row (sh.output, [ get_origin `Output ])
+    :: !active_constraints
+
 let unsafe_reinitialize () =
   update_uid := 0;
   state := Row.empty_env;
@@ -1473,7 +1490,6 @@ let%debug4_sexp propagate_shapes (update_step : update_step) : unit =
   iter_shapes update_step ~f:(apply_env_t !state);
   let _, ineqs = get_inequalities update_step in
   active_update_steps := update_step :: !active_update_steps;
-  let _debug_new_active_update_steps : update_step list = !active_update_steps in
   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
diff --git a/tensor/shape.mli b/tensor/shape.mli
@@ -176,6 +176,11 @@ val unsafe_reinitialize : unit -> unit
 
 (** {2 Internal-ish API.} *)
 
+val set_terminal : t -> unit
+(** Mark the shape as terminal, so that its rows can be closed to Least Upper Bounds (LUBs). This
+    function is only intended for parameters shapes, which would otherwise not be terminal because
+    of the initialization expressions of the parameters. *)
+
 (** How to propagate shape updates and do the last update of [Tensor.t.shape] when finalizing the
     tensor. Axes are broadcast-expanded on a bottom-up update to fit the incoming shape. *)
 type logic =
diff --git a/tensor/tensor.ml b/tensor/tensor.ml
@@ -589,6 +589,7 @@ let%debug7_sexp param ~t (name : string) ?(more_label = []) ?input_dims ?output_
      update computations. *)
   let g = (Option.value_exn ~here:[%here] t.diff).grad in
   Tn.update_memory_mode g Never_virtual 26;
+  Shape.set_terminal t.shape;
   remove_fwd_root t;
   { t with params = Set.singleton (module T) t }
 
diff --git a/test/einsum/test_einsum_capture.expected b/test/einsum/test_einsum_capture.expected
@@ -73,7 +73,7 @@ Test 1 - Matrix multiply with constraint-driven shapes:
   Captured dimensions: i=3, j=4, k=5
 
 Test 2 - Chain operations with constraint propagation:
-  base inferred shape: 1,1
+  base inferred shape: 6,8
   transposed inferred shape: 8,6
   multiplied inferred shape: 8,10
   final inferred shape: 6,10
