* 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 dimension variables in terminal shapes that do not have a LUB by dim 1, **except** when the variable has `has_uniq_constr_unless` set (indicating it's in the numerator of a `Total_elems` constraint) and none of the denominator variables are also prevented from guessing -- this prevents premature guessing that would make `Total_elems` constraints unsatisfiable. 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, or terminal ones, 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 7 sets all dim variables remaining in updated shapes to the lower bound if they have any, otherwise to dimension-1.

Let's explain the shape inference functions.

[%log6 "terminal row: keeping", (_r : row), "as", (r1 : row)];

Terminal_row (is_param, r1, origin) :: term_dims ())

match (lub, constr) with

| Some (Dim { d; _ } as lub), At_least_dim d2 when d2 > d ->

raise

@@ Shape_error

( [%string "dereferenced at dimension %{d2#Int}, higher than use site"],

[ Dim_mismatch [ lub; Var v ] ] )

| Some _, At_least_dim 0 ->

Some (Dim_eq { d1 = Var v; d2 = get_dim ~d ~proj_id:58 (); origin })

| _ -> None

let%track5_sexp process_shape_row ~(stage : stage) origin env ({ dims; bcast; prov } as r : row) :

("You forgot to specify the hidden dimension(s) 5", [ Row_mismatch [ r ] ])

| Some (Bounds_dim { lub; constr; has_uniq_constr_unless; _ })

when is_stage4_up stage && can_guess_dim_to_one env has_uniq_constr_unless ->

| Some (Solved_dim _) -> assert false

| Some (Bounds_dim { has_uniq_constr_unless; _ })

when final && can_guess_dim_to_one env has_uniq_constr_unless ->

let dim_eqs = process_dims beg_dims @ process_dims dims in

let r1 : row = row_of_var v prov in

match find_row env.row_env v with

| Some (Bounds_row { constr = Unconstrained; _ }) when not final ->

(Shape_row (r, origin) :: dim_eqs, env)

| Some (Bounds_row { constr = Unconstrained; _ }) when final ->

Output shape:

((batch

((dims

(bcast Broadcastable) (prov (((sh_id 64) (kind Batch))))))

(input ((dims ()) (bcast Broadcastable) (prov (((sh_id 64) (kind Input))))))

(output

(bcast Broadcastable) (prov (((sh_id 64) (kind Output))))))

(batch_padding ()) (input_padding ()) (output_padding ()) (id 64)

(debug_name output))

Output shape:

((batch

((dims

(bcast Broadcastable) (prov (((sh_id 154) (kind Batch))))))

(input

((dims ()) (bcast Broadcastable) (prov (((sh_id 154) (kind Input))))))

(output

(bcast Broadcastable) (prov (((sh_id 154) (kind Output))))))

(batch_padding ()) (input_padding ()) (output_padding ()) (id 154)

(debug_name layer_norm))

Retrieving commandline, environment, or config file variable ocannl_log_level

Found 0, in the config file

((batch

((dims

(input

(output

(debug_name transformer))