let ctx = Backend.make_context stream in

let open Operation.At in

CDSL.virtualize_settings.enable_device_only <- false;

let size = 50 in

let xs = Array.init size ~f:Float.(fun i -> (of_int i / 10.) + 0.1) in

let x_flat =

let rec apply_rows_constraint ~stage (rows : row list) (constr : row_constraint) (env : environment)

: constraint_ list * environment =

match rows_to_row_or_vars rows with

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

match constr with

| Exact dims when List.length dims < List.length all_dims ->

)

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

if is_unconstrained constr then ([], env)

else

let reduce constr ~beg_dims ~dims =

List.iter subr ~f:(fun subr ->

ineqs := Row_ineq { subr = row_of_var subr value.id; cur = r2 } :: !ineqs);

Option.iter lub ~f:(fun lub -> ineqs := Row_ineq { cur = lub; subr = r2 } :: !ineqs);

ineqs := extras @ !ineqs;

env)

else env

[]

let r_dims r =

match r.bcast with Broadcastable -> r.dims | Row_var { beg_dims; _ } -> beg_dims @ r.dims

(env : environment) : constraint_ list =

match rows_to_row_or_vars rows with

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

let rev_row_vars = List.rev row_vars in

match

| _, { kind = `Output; _ } -> true

| _ -> false) )

with

let other_vars = List.filteri rev_row_vars ~f:(fun i _ -> i <> idx) in

let other_vars = List.map other_vars ~f:(fun (v, id) -> row_of_var v id) in

let other_eqs =

| { bcast = Row_var { v = v'; _ }; _ } as r when equal_row_var v' v -> r

| r -> { r with dims = r_dims r; bcast = Broadcastable })

in

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

match r with

| { bcast = Broadcastable; _ } ->

(* The environment is unchanged, as apply_row_constraint would update only the constr. *)

List.concat_map ineqs ~f:(function

| ineq -> [ ineq ])

| { bcast = Row_var { v; beg_dims }; dims; id } -> (

let r1 = row_of_var v id in

match reduce_row_constraint constr ~beg_dims ~dims with

| Total_elems { numerator; divided_by } -> (

match (numerator, divided_by, lub) with

no_further_axes

:: List.map vs ~f:(fun v ->

let d2 = get_dim ~d:1 () in

Dim_eq { d1 = Var v; d2 })

let dim = get_dim ~d () in

[ Row_eq { r1; r2 = { dims = [ dim ]; bcast = Broadcastable; id } } ]

let dim = get_dim ~d () in

[ Row_eq { r1; r2 = { dims = [ dim ]; bcast = Broadcastable; id } } ]

| Num_elems _, [], Some lub ->

List.concat_map ineqs ~f:(function

| Rows_constr { r = rows; constr } ->

| ineq -> [ ineq ])

no_further_axes :: [ Dim_eq { d1 = Var v; d2 = get_dim ~d () } ]

(* The row variable should have coeff * var elements total *)

(* We can't determine the exact shape without knowing var *)

| Exact dims -> [ Row_eq { r1; r2 = { dims; bcast = Broadcastable; id } } ]

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

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

match bcast with

| Row_var { v; beg_dims } -> (

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

let r1 : row = row_of_var v id in

let no_further_axes = Row_eq { r1; r2 = { dims = []; bcast = Broadcastable; id } } in

match Map.find env.row_env v with

[%log6 "terminal row: closing", (_r : row)];

no_further_axes :: term_dims ()

| Some (Bounds_row { lub = None; constr; _ })

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

let ineqs =

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

| Some (Solved_row _) -> assert false

| Some (Bounds_row { lub = Some lub; _ }) when is_stage3_up stage ->

Row_eq { r1; r2 = lub } :: term_dims ()

| _ ->

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

Terminal_row r1 :: term_dims ())

| Bounds_row { lub; constr; _ } ->

(* TODO: should we store the id somewhere? *)

let id = phantom_row_id in

| _ -> []

in

let finalizing_entries : constraint_ list =

The constraints are solved by: unification of the equation constraints, unification-like simplification of the inequality constraints, propagation of the complex constraints. Simplification of an inequality, and constraint propagation, can generate more constraints, so we need to be careful to keep it terminating. The solution proceeds in stages.

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

* Stage 7 sets all dim resp. row variables remaining in updated shapes to dimension-1 resp. no-further-axes.