To separate concerns, OCANNL is split into the `arrayjit` library, responsible for compilation of high-level n-D array operation sequences (`Assignments.comp`) via backends such as sync_cc, metal and cuda, and the main `ocannl` library, responsible for deriving the operations computing the forward propagation and backpropagation from tensor expressions. In particular, `arrayjit` contains `Indexing`, which represents complex indexing into arrays, and the main library `ocannl` has `Row` and `Shape` modules, which do the most "heavy-lifting" in the translation from concise tensor expressions to sequences of assignments.

The bulk of the projections inference happens alongside shape inference, with the projections-relevant information stored in auxiliary fields -- this prevents subtle bugs where projection semantics deviates from shape semantics, and will simplify adding new shape/projection inference features. Shape inference happens during `propagate_shapes` calls, and then again in a `finish_inference` call, which is triggered whenever the dimensions or projections are required (i.e. typically by jitting). Finally, the projections are reconstructed in `derive_projections`. It would seem `derive_projections` could reuse the already-computed solutions constraints. But we face a problem: we must prevent contaminating projections across different operations. To illustrate: we conclude the dimensions of two axes are the same because they are reduced together in another operation -- this should not force the axes to share a projection in the processed operation. To prevent the contamination, in each `derive_projections` call, we freshen the projection ids in the (inferred) shapes, and regenerate and re-solve the constraints with the fresh projection ids.

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. Currently there are 8 stages, with a fractional stage coming from splitting an earlier design.

@@ Shape_error

("dimension comparison for axis: different labels", [ Dim_mismatch [ cur; subr ] ])

| Dim { d = d1; _ }, Dim { d = d2; _ } when d1 = d2 -> ([], env)

| Conv_input _, _ | _, Conv_input _ -> ([ Dim_eq { d1 = subr; d2 = cur; origin } ], env)

| Var cur_v, Var subr_v -> (

match (find_dim env.dim_env cur_v, find_dim env.dim_env subr_v) with

List.map2_exn (take_from_end r_cur.dims lub_len) (take_from_end lub2.dims lub_len)

~f:(fun d1 d2 ->

match (d1, d2) with

| Dim { d = d1; _ }, Dim { d = d2; _ } when d1 <> d2 -> get_dim ~d:1 ~proj_id:48 ()

| Conv_input { stride; output = Dim s; _ }, Dim s'

| Dim s', Conv_input { stride; output = Dim s; _ }

when !use_padding && stride * s.d <> s'.d ->

let back_axes, front_axes =

Map.to_alist axes

|> List.partition_map ~f:(fun ({ AxisKey.from_end; pos = i; _ }, v) ->

in

let back_size = List.fold back_axes ~init:0 ~f:(fun accu (i, _) -> max i accu) in

let front_size = List.fold front_axes ~init:0 ~f:(fun accu (i, _) -> max i accu) in

( proj_axis_env,

(Option.to_list static_range

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

@ [

Row_eq { r1 = expanded_batch; r2 = sh.batch; origin = get_origin `Batch };

Row_eq { r1 = cur_sh.input; r2 = sh.input; origin = get_origin `Input };

let make ?batch_dims ?input_dims ?output_dims ?batch_axes ?input_axes ?output_axes

?(deduced = Not_constrained) ~debug_name ~id () =

let open Row in

let make_dims kind ds =

in

let make_axes kind ds =

{

let dims = (row_of_kind kind sh).dims in

String.concat ~sep:","

@@ List.mapi dims ~f:(fun i d ->

in

let batch_dims = dims_to_string `Batch in

let input_dims = dims_to_string `Input in