def output_eq

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as Par
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

985# The writes of every cell, replayed on the owned storage `created` makes for
986# the target, give the owned storage the parallel build fills.
987def output_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
988  +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>},
989  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +dc: Nat, +fc: F32,
990  +source: Par.Source<shape, fill, ops, h, dc, fc>, +target: Shape.Shape, +facts: Facts<plan, shape, U32.to_nat(Shape.count(target))>,
991  +d: Nat, +fits: {Storage.bounds(d, U32.to_nat(Shape.count(target))) == True{} : Bool},
992  -f: {Storage.bounds(d, U32.to_nat(Shape.count(target))) == True{} : Bool}) ->
993  {StorageLaws.replay(writes(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, U32.to_nat(Shape.count(target))),
994    Storage.Owned{d, U32.to_nat(Shape.count(target)), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}) ==
995    Storage.Owned{d, U32.to_nat(Shape.count(target)),
996      Reduce.build(d, parallel_splits(plan, d, U32.to_nat(Shape.count(target))), rows, Shape.extents(Shape.output_shape(plan)),
997        U32.to_nat(Shape.count(target)), 0n, Api.arr(h, dc, fc)),
998      Reduce.build_balanced(d, parallel_splits(plan, d, U32.to_nat(Shape.count(target))), rows, Shape.extents(Shape.output_shape(plan)),
999        U32.to_nat(Shape.count(target)), 0n, Api.arr(h, dc, fc)), f} : Storage.Storage}:
1000  +n = U32.to_nat(Shape.count(target))
1001  +map : PlanMap(plan) = Shape.Indexed{rows, rows_exact}
1002  +splits = parallel_splits(plan, d, n)
1003  +extents = Shape.extents(Shape.output_shape(plan))
1004  Equal.trans(Storage.Storage,
1005    StorageLaws.replay(writes(plan, map, shape, fill, ops, n), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
1006    Storage.Owned{d, n, out_spec(plan, map, shape, fill, ops, n, d, 0n), out_balanced(plan, map, shape, fill, ops, n, d, 0n), f},
1007    Storage.Owned{d, n, Reduce.build(d, splits, rows, extents, n, 0n, Api.arr(h, dc, fc)), Reduce.build_balanced(d, splits, rows, extents, n, 0n, Api.arr(h, dc, fc)), f},
1008    chain(plan, map, shape, fill, ops, d, n, f, fits, Shape.count(target), Conversion.le_reflexive(n), n, Conversion.le_reflexive(n)),
1009    Owned.owned_eq(d, n, out_spec(plan, map, shape, fill, ops, n, d, 0n), Reduce.build(d, splits, rows, extents, n, 0n, Api.arr(h, dc, fc)),
1010      Equal.sym(Array<F32>, Reduce.build(d, splits, rows, extents, n, 0n, Api.arr(h, dc, fc)), out_spec(plan, map, shape, fill, ops, n, d, 0n),
1011        build_eq(plan, rows, rows_exact, shape, fill, ops, h, dc, fc, n, source, facts, d, splits, 0n)),
1012      out_balanced(plan, map, shape, fill, ops, n, d, 0n), f)(Reduce.build_balanced(d, splits, rows, extents, n, 0n, Api.arr(h, dc, fc))))