def build_cells

bendtorch-core@0.1.0.0 · src/tensor/matmul_parallel_proof.bend:238

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic

238# Whatever the split, every cell of the parallel build holds what CellIs
239# names: the generic law of `Parallel.build`, with the matmul cell.
240def build_cells(+depth: Nat, +splits: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +inner: Nat, +columns: Nat,
241  +count: Nat, +base: Nat, la: Array<F32>, ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>},
242  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
243  Parallel.Fact<Cells(depth, dl, dr, A, B, inner, columns, count, base, Matrix.build(depth, splits, inner, columns, count, base, la, ra))>:
244  -split = {\{Parallel.Fact: q => q} : Parallel.Fact<({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>} &
245    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>})> ->
246    {Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>} &
247    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}}(
248    Parallel.split_equal(dl, dr, A, B, left_valid, right_valid))
249  -cells = {\{Generic.Fact: q => q} :
250    Generic.Fact<Generic.Cells(~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.cell_at, Matrix.dims(Matrix.Dims{inner, columns}), (A, B), depth, count, base,
251      Matrix.build(depth, splits, inner, columns, count, base, la, ra))> ->
252    Generic.Cells(~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.cell_at, Matrix.dims(Matrix.Dims{inner, columns}), (A, B), depth, count, base,
253      Matrix.build(depth, splits, inner, columns, count, base, la, ra))}(
254    Generic.build_cells(~Matrix.Dims, ~Matrix.Dims, ~Array<F32>, ~Array<F32>, ~Matrix.dims, ~Elementwise.split, ~Matrix.cell_at,
255      Matrix.Dims{inner, columns}, (A, B), depth, splits, count, base, (la, ra), Parallel.pair_equal(la, ra, A, B, el, er),
256      at => inside => {\{Parallel.Fact: q => q} :
257        Parallel.Fact<({Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, at, A, B, inside)) == (A, B) : Array<F32> & Array<F32>})> ->
258        {Pair.snd(F32, Array<F32> & Array<F32>, cell_of(inner, columns, at, A, B, inside)) == (A, B) : Array<F32> & Array<F32>}}(
259        cell_keeps(dl, dr, A, B, inner, columns, at, inside, left_valid, right_valid)),
260      Pair.fst({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>},
261        {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}, split),
262      Pair.snd({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>},
263        {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}, split)))
264  Parallel.Fact{path => {\{Parallel.Fact: q => q} :
265    Parallel.Fact<CellSpec(dl, dr, A, B, inner, columns, Nat.add(Model.offset(depth, path), base), Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
266      Model.observe(F32, depth, Matrix.build(depth, splits, inner, columns, count, base, la, ra), path))> ->
267    CellSpec(dl, dr, A, B, inner, columns, Nat.add(Model.offset(depth, path), base), Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
268      Model.observe(F32, depth, Matrix.build(depth, splits, inner, columns, count, base, la, ra), path))}(
269    cell_spec(dl, dr, A, B, inner, columns, Nat.add(Model.offset(depth, path), base), Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
270      Model.observe(F32, depth, Matrix.build(depth, splits, inner, columns, count, base, la, ra), path), cells(path), left_valid, right_valid))}