| Some (Bounds_row { constr = Total_elems _; _ }) -> assert false

| _ -> elim_var :: elim_dims)

let%debug4_sexp solve_inequalities ~(stage : stage) (ineqs : constraint_ list) (env : environment) :

constraint_ list * environment =

| Some i -> i

| None -> unknown_projection proj_id d)

| Solved idx -> idx

let output_idx = loop output in

let symbols = ref [] in

let offset = ref 0 in

(* Expand output index - multiply by stride *)

(match output_idx with

| Idx.Fixed_idx i -> offset := !offset + (stride * i)

| Idx.Iterator s -> symbols := (stride, s) :: !symbols

| Idx.Affine { symbols = output_syms; offset = output_offset } ->

symbols := List.map output_syms ~f:(fun (c, s) -> (stride * c, s)) @ !symbols;

offset := !offset + (stride * output_offset));

(* Process solved kernel terms *)

List.iter solved_kernel ~f:(fun (dilation, idx) ->

match idx with

| Idx.Affine { symbols = kernel_syms; offset = kernel_offset } ->

symbols := List.map kernel_syms ~f:(fun (c, s) -> (dilation * c, s)) @ !symbols;

offset := !offset + (dilation * kernel_offset));

(* Process unsolved kernel terms *)

List.iter unsolved_kernel ~f:(fun (dilation, v) ->

let idx = loop (Var v) in

| Idx.Affine { symbols = kernel_syms; offset = kernel_offset } ->

symbols := List.map kernel_syms ~f:(fun (c, s) -> (dilation * c, s)) @ !symbols;

offset := !offset + (dilation * kernel_offset));

(* Add padding if use_padding is true *)

if !use_padding then

match output with

| Proj (proj_id, _) ->

Utils.union_find ~equal:Proj_id.equal proj_env.proj_classes ~key:proj_id ~rank:0

in

Option.value (Map.find proj_env.proj_to_padding repr) ~default:0

in

| _ -> !offset

else !offset

in

(* Combine and simplify symbols *)

!symbols

|> List.filter ~f:(fun (c, _) -> c <> 0)

|> List.sort ~compare:(fun (_, s1) (_, s2) -> Idx.compare_symbol s1 s2)

|> List.group ~break:(fun (_, s1) (_, s2) -> not (Idx.equal_symbol s1 s2))

|> List.map ~f:(fun group ->

|> List.filter ~f:(fun (c, _) -> c <> 0)

in

| [] -> Idx.Fixed_idx offset

| _ -> Idx.Affine { symbols; offset })

| Var v when Hashtbl.mem proj_env.v_env v -> loop (Hashtbl.find_exn proj_env.v_env v)

| Var v ->

| Dim s -> Proj (Proj_id.fresh (), s)

| Conv_input { stride; output; solved_kernel; unsolved_kernel } ->

let output_proj = dim_to_proj output in

match solved_kernel with

(* Convert solved_dim to axis_index *)

if Idx.iterated sd.d then

match sd.proj_id with

in

| Some idx -> idx

| None -> Idx.Fixed_idx 0)

| None -> Idx.Fixed_idx 0

else Idx.Fixed_idx 0

in

| None -> []

in

in

in

loop

(* Process p_conv_input to populate projs and compute padding *)

let proj_to_padding = ref @@ Map.empty (module Proj_id) in

(* Helper to compute padding from Conv_input projection *)

let compute_padding_from_proj = function

| Conv_input { unsolved_kernel; _ } ->

(* Padding is the maximum dilation factor among kernel terms *)

List.fold unsolved_kernel ~init:0 ~f:(fun acc (dilation, _) -> Int.max acc dilation)

| _ -> 0

in

(* Process postponed Conv_input equations *)

List.iter !p_conv_input ~f:(fun (p, conv_input) ->

let repr, _ = Utils.union_find ~equal:Proj_id.equal !proj_classes ~key:p ~rank:0 in

(* Substitute variables in conv_input to get resolved projection *)

let rec substitute_vars_in_proj = function

| Conv_input { stride; output; solved_kernel; unsolved_kernel } ->

let output' = substitute_vars_in_proj output in

Conv_input { stride; output = output'; solved_kernel; unsolved_kernel }

| Var v as proj -> (

| proj -> proj

in

let resolved_conv_input = substitute_vars_in_proj conv_input in

(* Compute padding if use_padding is true *)

(* Try to get index for Conv_input - this creates a temporary proj_env *)

try

let idx = get_proj_index temp_proj_env resolved_conv_input in

Utils.mref_add projs ~key:repr ~data:idx ~or_:(fun idx2 ->

if not @@ Idx.equal_axis_index idx idx2 then

raise

@@ Shape_error

(* Verify postponed equations *)

List.iter !verify_when_solved1 ~f:(fun (idx, conv_input) ->

try

let conv_idx = get_proj_index temp_proj_env conv_input in

if not @@ Idx.equal_axis_index idx conv_idx then

raise

@@ Shape_error

List.iter !verify_when_solved2 ~f:(fun (conv_input1, conv_input2) ->

try

let idx1 = get_proj_index temp_proj_env conv_input1 in

let idx2 = get_proj_index temp_proj_env conv_input2 in

if not @@ Idx.equal_axis_index idx1 idx2 then

raise

@@ Shape_error

("Cannot unify two Conv_input projections", [ Index_mismatch [ idx1; idx2 ] ])

{

v_env;

The shape and projection inference handles `Conv_input` terms differently depending on `use_padding`. If `use_padding` is false, the impact of convolution kernels is incorporated additively during shape inference, and there's nothing more to do during projection inference. If `use_padding` is true, convolution kernels don't contribute during shape inference, and padding is computed during projection inference, keyed by `proj_id`. Padding is maximal size (width) of dilated kernels as encountered in `Dim` - `Conv_input` constraints and is propagated in either direction, although in practice for CNNs `Conv_input` should only appear as `subr` of `Dim_ineq`.

The actual shape inference combines row polymorphism with (nominal) subtyping, as known in the type inference literature. The subtyping stems merely from the fact that a dimension-1 axis can be used in the context of any dimension due to per-axis broadcasting. Row polymorphism stems from broadcasting to more axes: for example, when unifying an unknown (shape) row with a known one, we cannot assume that the unknown row will have just the axes of the known one, because maybe the known row is meant to be broadcasted here to more axes. The combination of row polymorphism with nominal subtyping means that the constraints we are solving are inequalities, both inequalities between rows (the `Row.t` type, i.e. the `row` type above), and between axes/dimensions (the `Row.dim` type). We maintain the inequality ordering between variables in the environment to compute the transitive closure during simplification. We also maintain a least upper bound on the solution.

* `solve_proj_equations` unifies the projection equations, using union-find to maintain a representative for equal projections. Projections that already have an `axis_index` are `non_product` (not to be iterated over). The remaining projections have a `product_dim`, and get a fresh iterator.

* `get_dim_index` gets an `axis_index` for a `dim` based on the representative of its `proj_id`; and `Fixed_idx 0` for dim=1.

Other important functions in the `Shape` module.