def build_cells

bendtorch-core@0.1.0.0 · src/tensor/gather_PROOF.bend:336

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/tensor/gather.bend as Gather
import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic

336# Every slot of the forward build holds what its cell computes from (A, I).
337# The Fact holds the cells eta-expanded, as slice_PROOF.bend's build_cells.
338def build_cells(+depth: Nat, +splits: Nat, +plan: Gather.Plan, +cells: Nat, +base: Nat, +da: Nat, -A: Array<F32>, a: Array<F32>,
339  -same: {a == A : Array<F32>}, -valid: Model.Balanced(F32, da, A), +di: Nat, -I: Array<U32>, index: Array<U32>,
340  -same_index: {index == I : Array<U32>}, -index_valid: Model.Balanced(U32, di, I)) ->
341  Generic.Fact<Generic.Cells(~Gather.Plan, ~Array<F32>, ~Array<U32>, ~Gather.cell, plan, (A, I), depth, cells, base, Gather.build(depth, splits, plan, cells, base, a, index))>:
342  Generic.Fact{path =>
343    {\{Generic.Fact: q => q} : Generic.Fact<Generic.Cells(~Gather.Plan, ~Array<F32>, ~Array<U32>, ~Gather.cell, plan, (A, I), depth, cells, base, Gather.build(depth, splits, plan, cells, base, a, index))> -> Generic.Cells(~Gather.Plan, ~Array<F32>, ~Array<U32>, ~Gather.cell, plan, (A, I), depth, cells, base, Gather.build(depth, splits, plan, cells, base, a, index))}(
344      Generic.build_cells(~Gather.Plan, ~Gather.Plan, ~Array<F32>, ~Array<U32>, ~Gather.same, ~Gather.split, ~Gather.cell,
345        plan, (A, I), depth, splits, cells, base, (a, index), both(a, index, A, I, same, same_index),
346        at => inside => {\{Generic.Fact: q => q} :
347          Generic.Fact<{Pair.snd(F32, Array<F32> & Array<U32>, Gather.cell(plan, at, (A, I), inside)) == (A, I) : Array<F32> & Array<U32>}> ->
348          {Pair.snd(F32, Array<F32> & Array<U32>, Gather.cell(plan, at, (A, I), inside)) == (A, I) : Array<F32> & Array<U32>}}(
349          cell_keeps(plan, at, da, A, di, I, inside, valid, index_valid)),
350        Pair.fst({Pair.fst(Array<F32> & Array<U32>, Array<F32> & Array<U32>, Gather.split((A, I))) == (A, I) : Array<F32> & Array<U32>}, {Pair.snd(Array<F32> & Array<U32>, Array<F32> & Array<U32>, Gather.split((A, I))) == (A, I) : Array<F32> & Array<U32>},
351          {\{Generic.Fact: q => q} : Generic.Fact<SplitFacts(A, I)> -> SplitFacts(A, I)}(split_facts(da, A, di, I, valid, index_valid))),
352        Pair.snd({Pair.fst(Array<F32> & Array<U32>, Array<F32> & Array<U32>, Gather.split((A, I))) == (A, I) : Array<F32> & Array<U32>}, {Pair.snd(Array<F32> & Array<U32>, Array<F32> & Array<U32>, Gather.split((A, I))) == (A, I) : Array<F32> & Array<U32>},
353          {\{Generic.Fact: q => q} : Generic.Fact<SplitFacts(A, I)> -> SplitFacts(A, I)}(split_facts(da, A, di, I, valid, index_valid)))))(path)}