def cell_eq

bendtorch-core@0.1.0.0 · src/autograd/reduce_replayed_proof.bend:737

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as Par
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
import bendtorch-core@0.1.0.0/src/autograd/reduce_parallel_proof.bend as ReduceParallel

737def cell_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
738  +rows_exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>},
739  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32, +count: Nat,
740  +source: Par.Source<shape, fill, ops, h, d, f>, +facts: Facts<plan, shape, count>, +base: Nat, inside: Bool,
741  +decided: {inside == Nat.is_lt(base, count) : Bool}) ->
742  TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 0n, base,
743    Parallel.leaf(~(Array<F32> & Unit), Reduce.keep(Reduce.cell(rows, Shape.extents(Shape.output_shape(plan)), base, Api.arr(h, d, f), inside), Unit{}))):
744  match inside:
745    case False{}:
746      (%decided : {ALeaf{0.0} == ALeaf{out_leaf(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, base, _)} : Array<F32>}
747       {==}, {==})
748    case True{}:
749      +output = Shape.extents(Shape.output_shape(plan))
750      +word = rfold(shape, fill, ops, ReduceParallel.visited_of(rows, output, base), 0.0)
751      %Equal.sym(F32 & Array<F32>, Reduce.fiber_of(rows, output, base, Api.arr(h, d, f)), (word, Api.arr(h, d, f)),
752        Equal.trans(F32 & Array<F32>, Reduce.fiber_of(rows, output, base, Api.arr(h, d, f)),
753          Reduce.fiber(rows, output, Reduce.first_span(rows, output), base, 0n, Reduce.first_digit(rows, base), 0.0, Api.arr(h, d, f)),
754          (word, Api.arr(h, d, f)),
755          fast_arr(rows, output, Reduce.first_span(rows, output), Reduce.single_below(rows, output),
756            ReduceParallel.single_ok(rows, output, Reduce.single_below(rows, output), {==}), base, 0n, Reduce.first_digit(rows, base), 0.0, h, d, f),
757          fiber_eq(rows, output, Reduce.first_span(rows, output), base, 0n, Reduce.first_digit(rows, base), 0.0, shape, fill, ops, h, d, f, source,
758            walk_below(plan, rows, rows_exact, shape, count, facts, base, Equal.sym(Bool, True{}, Nat.is_lt(base, count), decided))))) :
759        TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 0n, base, Parallel.leaf(~(Array<F32> & Unit), Reduce.keep(_, Unit{})))
760      (%decided : {ALeaf{word} == ALeaf{out_leaf(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, base, _)} : Array<F32>}
761       {==}, {==})