def fork_eq

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
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/binary_build_proof.bend as Build
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

849def fork_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
850  +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>},
851  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32, +count: Nat,
852  +lower: Nat, +rest: Nat, +base: Nat,
853  low: {Reduce.build(lower, rest, rows, Shape.extents(Shape.output_shape(plan)), count, base, Api.arr(h, d, f)) ==
854    out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, lower, base) : Array<F32>},
855  high: {Reduce.build(lower, rest, rows, Shape.extents(Shape.output_shape(plan)), count, Nat.add(Model.capacity(lower), base), Api.arr(h, d, f)) ==
856    out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, lower, Nat.add(Model.capacity(lower), base)) : Array<F32>},
857  pair: Array<F32> & Array<F32>, +exact: {pair == Array.clone(F32, Api.arr(h, d, f)) : Array<F32> & Array<F32>}) ->
858  {Parallel.fork(~Array<F32>, ~Unit, Reduce.halves(pair),
859    l => r => Parallel.build(~Reduce.Fibers, ~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.fibers, ~Reduce.split, ~Reduce.cell_at, Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}, lower, rest, count, base, (l, r)),
860    l => r => Parallel.build(~Reduce.Fibers, ~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.fibers, ~Reduce.split, ~Reduce.cell_at, Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}, lower, rest, count, Nat.add(Model.capacity(lower), base), (l, r))) ==
861    out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 1n+lower, base) : Array<F32>}:
862  (first, second) = pair
863  +output = Shape.extents(Shape.output_shape(plan))
864  +upper = Nat.add(Model.capacity(lower), base)
865  %Equal.sym(Array<F32>, first, Api.arr(h, d, f), Build.clone_first_eq(d, h, f, (first, second), exact)) :
866    {ANode{Reduce.build(lower, rest, rows, output, count, base, _), Reduce.build(lower, rest, rows, output, count, upper, second)} == out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 1n+lower, base) : Array<F32>}
867  %Equal.sym(Array<F32>, second, Api.arr(h, d, f), Build.clone_second_eq(d, h, f, (first, second), exact)) :
868    {ANode{Reduce.build(lower, rest, rows, output, count, base, Api.arr(h, d, f)), Reduce.build(lower, rest, rows, output, count, upper, _)} == out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 1n+lower, base) : Array<F32>}
869  %Equal.sym(Array<F32>, Reduce.build(lower, rest, rows, output, count, base, Api.arr(h, d, f)),
870    out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, lower, base), low) :
871    {ANode{_, Reduce.build(lower, rest, rows, output, count, upper, Api.arr(h, d, f))} == out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 1n+lower, base) : Array<F32>}
872  %Equal.sym(Array<F32>, Reduce.build(lower, rest, rows, output, count, upper, Api.arr(h, d, f)),
873    out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, lower, upper), high) :
874    {ANode{out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, lower, base), _} == out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 1n+lower, base) : Array<F32>}
875  {==}