| Rows_constr of { r : t list; constr : row_constraint }

| Terminal_dim of dim

| Terminal_row of t

[@@deriving compare, equal, sexp_of, variants]

type stage = Stage1 | Stage2 | Stage3 | Stage4 | Stage5 | Stage6 | Stage7

| Row_eq { r1; r2 } ->

let more_ineqs, env = unify_row ~stage (r1, r2) env in

(more_ineqs @ ineqs, env)

(ineq :: ineqs, env)

in

let unify_suffix init dims1 dims2 len =

else

match rows_to_row_or_vars rows with

| Either.First single_row ->

| Either.Second (_all_dims, row_vars) -> (

let rev_row_vars = List.rev row_vars in

match

(other_eqs @ ineqs, env)

| _ -> ([ Rows_constr { r = rows; constr } ], env))

constraint_ list * environment =

let keep_constr () =

let ineqs, env = apply_row_constraint ~depth stage r constr env in

let dim = get_dim ~d () in

([ Row_eq { r1; r2 = { dims = [ dim ]; bcast = Broadcastable; id } } ], env)

| Num_elems _, [], Some lub ->

List.fold ineqs ~init:([], env) ~f:(fun (ineqs, env) ineq ->

match ineq with

| Rows_constr { r = rows; constr } ->

let%track5_sexp close_row_terminal ~(stage : stage) (env : environment)

({ dims; bcast; id } as _r : row) : constraint_ list =

let suffix () = List.map dims ~f:(fun d -> Terminal_dim d) in

match bcast with

| Broadcastable -> if is_stage6_up stage then [] else suffix ()

| Row_var { v; beg_dims } -> (

when is_stage2_up stage && not (equal_row_constraint constr Unconstrained) ->

let ineqs, _env =

(* This is the constraint on the row variable, not on the original row. *)

with Shape_error (s, trace) -> raise @@ Shape_error (s, Row_mismatch [ r1 ] :: trace)

in

(* FIXME: at which stage should we drop the terminal row? *)

let%debug5_sexp eliminate_dim_entry ~final v ~lub constr =

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 lub, At_least_dim _ -> Some (Dim_eq { d1 = Var v; d2 = lub })

in

match bcast with

| Row_var { v; beg_dims } -> (

match Map.find env.row_env v with

in

let empty_env = { dim_env = Map.empty (module Dim_var); row_env = Map.empty (module Row_var) }

| Terminal_row r ->

let more_ineqs = close_row_terminal ~stage env @@ subst_row env r in

(more_ineqs @ ineqs, env)

in

let ineqs', env = List.fold ineqs ~init:([], env) ~f in

let ineqs' = List.rev ineqs' in

then (ineqs', env)

else solve ineqs' env

in

let rec row_to_labels env =

let rec f = function

| eq -> [ eq ])

| Terminal_dim d -> [ Iterated (to_proj d) ]

| Terminal_row { dims; _ } -> List.map ~f:(fun d -> Iterated (to_proj d)) dims

| Rows_constr

{ r; constr = Total_elems { numerator = Strided_var { coeff; var; denom }; divided_by } }

-> (

| Terminal_dim of dim

| Terminal_row of t

[@@deriving compare, equal, sexp_of, variants]

type stage = Stage1 | Stage2 | Stage3 | Stage4 | Stage5 | Stage6 | Stage7

val subst_row : environment -> t -> t

val unify_row : stage:stage -> t * t -> environment -> constraint_ list * environment

val empty_env : environment

val solve_inequalities :

stage:stage -> constraint_ list -> environment -> constraint_ list * environment

sh.input <- Row.subst_row env sh.input;

sh.output <- Row.subst_row env sh.output

let%debug4_sexp propagate_shapes (update_step : update_step) : unit =

(* Allow the derivation of constraints to depend on the shapes (currently, only Batch_slice

let _, ineqs = get_inequalities update_step in

active_update_steps := update_step :: !active_update_steps;

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

state := env

let%debug4_sexp finish_inference (() : unit) : unit =

(* TODO: optimize to keep all needed information in unsolved, rather than starting with all

let unsolved, env = Row.solve_inequalities ~stage:Stage2 !active_constraints !state in

let unsolved, env = Row.solve_inequalities ~stage:Stage3 unsolved env in

let unsolved, env = Row.solve_inequalities ~stage:Stage4 unsolved env in

let unsolved, env = Row.solve_inequalities ~stage:Stage5 unsolved env in

let unsolved, env = Row.solve_inequalities ~stage:Stage6 unsolved env in

let unsolved, env = Row.solve_inequalities ~stage:Stage7 unsolved env in

assert (List.is_empty unsolved);

let _applied_update_steps : update_step list = !active_update_steps in

active_constraints := [];

active_update_steps := [];

* 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.

* 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.

Let's explain the shape inference functions.