src/autograd/reduce_fold_PROOF.bend
bendtorch-core@0.1.0.0 · 263 lines · 13 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as ShapeLaws
4import bendtorch-core@0.1.0.0/src/shape/PROOF.bend as ShapeProof
5import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
6import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
7import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
8import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
9import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
10import bendtorch-core@0.1.0.0/src/tensor/projection_LAWS.bend as Projection
11import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
12import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
13import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
14import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as Broadcast
15import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast_proof.bend as BroadcastProof
16import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
17import bendtorch-core@0.1.0.0/src/tensor/contracts/fold_records.bend as Records
18import bendtorch-core@0.1.0.0/src/tensor/contracts/fold_proof.bend as FoldProof
19import bendtorch-core@0.1.0.0/src/tensor/contracts/sum.bend as Sum
20import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
21import bendtorch-core@0.1.0.0/src/tensor/contracts/PROOF.bend as ContractsProof
22import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
23import bendtorch-core@0.1.0.0/src/autograd/reduce_fold.bend as ReduceFold
24import bendtorch-core@0.1.0.0/src/autograd/reduce_fold_LAWS.bend as Laws
25
26def offsets_walk(+plan: Shape.BroadcastPlan,
27 +map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>,
28 +position: Nat, +bound: {True{} == Nat.is_lt(position, Shape.elements(Shape.output_shape(plan))) : Bool}) ->
29 {Projection.offset_spec(plan, position) ==
30 Some{(U32.from_nat(Broadcast.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
31 Shape.extents(Shape.output_shape(plan)), map, position)),
32 U32.from_nat(Broadcast.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)),
33 Shape.extents(Shape.output_shape(plan)), map, position)))} : Maybe<&1, U32 & U32>}:
34 match map:
35 case Shape.Indexed{+rows, +exact}:
36 +left = Shape.extents(Shape.left_shape(plan))
37 +right = Shape.extents(Shape.right_shape(plan))
38 +output = Shape.extents(Shape.output_shape(plan))
39 +index = {Shape.FIndex{position, bound} : Shape.FlatIndex<Shape.output_shape(plan)>}
40 +slot = ShapeLaws.fin_of(Shape.numel(output), position)
41 +x = U32.from_nat(Shape.walk_left(rows, position))
42 +y = U32.from_nat(Shape.walk_right(rows, position))
43 %ShapeLaws.project_left_index_refines(plan, Shape.Indexed{rows, exact}, index) :
44 {Projection.pair_offsets(_, ShapeLaws.flatten_spec(right, ShapeLaws.project_right_coordinate(left, right, output, Shape.alignment(plan), slot))) ==
45 Some{(x, y)} : Maybe<&1, U32 & U32>}
46 %ShapeLaws.project_right_index_refines(plan, Shape.Indexed{rows, exact}, index) :
47 {Projection.pair_offsets(Some{x}, _) == Some{(x, y)} : Maybe<&1, U32 & U32>}
48 {==}
49 case Shape.Unindexed{zero}:
50 +output = Shape.output_shape(plan)
51 Empty.absurd({Projection.offset_spec(plan, position) == Some{(U32.from_nat(0n), U32.from_nat(0n))} : Maybe<&1, U32 & U32>},
52 Arith.lt_zero(position,
53 %Equal.trans(Nat, Shape.elements(output), Shape.numel(Shape.extents(output)), 0n, Shape.elements_exact(output),
54 Equal.sym(Nat, 0n, Shape.numel(Shape.extents(output)), zero)) : {Nat.is_lt(position, _) == True{} : Bool}
55 Equal.sym(Bool, True{}, Nat.is_lt(position, Shape.elements(output)), bound)))
56
57def length_snoc(+values: List<&2, F32>, +word: F32) ->
58 {List.length(&2, F32, List.append(&2, F32, values, [word])) == 1n+List.length(&2, F32, values) : Nat}:
59 match values:
60 case Nil{}:
61 {==}
62 case +value <> +rest:
63 Equal.cong(Nat, Nat, n => 1n+n, List.length(&2, F32, List.append(&2, F32, rest, [word])), 1n+List.length(&2, F32, rest),
64 length_snoc(rest, word))
65
66def words_length(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) ->
67 {List.length(&2, F32, Contracts.words(shape, fill, ops, count)) == count : Nat}:
68 match count:
69 case 0n:
70 {==}
71 case 1n+previous:
72 +before = Contracts.words(shape, fill, ops, previous)
73 Equal.trans(Nat, List.length(&2, F32, List.append(&2, F32, before, [Replayed.value(shape, fill, ops, U32.from_nat(previous))])),
74 1n+List.length(&2, F32, before), 1n+previous,
75 length_snoc(before, Replayed.value(shape, fill, ops, U32.from_nat(previous))),
76 Equal.cong(Nat, Nat, n => 1n+n, List.length(&2, F32, before), previous, words_length(shape, fill, ops, previous)))
77
78def word_at_front(+values: List<&2, F32>, +word: F32, +position: Nat,
79 below: {Nat.is_lt(position, List.length(&2, F32, values)) == True{} : Bool}) ->
80 {ReduceLaws.word_at(List.append(&2, F32, values, [word]), position) == ReduceLaws.word_at(values, position) : Maybe<&2, F32>}:
81 match values position:
82 case Nil{} _:
83 Empty.absurd({ReduceLaws.word_at([word], position) == None{} : Maybe<&2, F32>}, Arith.lt_zero(position, below))
84 case +value <> +rest 0n:
85 {==}
86 case +value <> +rest 1n+next:
87 word_at_front(rest, word, next, below)
88
89def word_at_last(+values: List<&2, F32>, +word: F32) ->
90 {ReduceLaws.word_at(List.append(&2, F32, values, [word]), List.length(&2, F32, values)) == Some{word} : Maybe<&2, F32>}:
91 match values:
92 case Nil{}:
93 {==}
94 case +value <> +rest:
95 word_at_last(rest, word)
96
97def words_step(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +previous: Nat, +position: Nat,
98 below: {Nat.is_lt(position, 1n+previous) == True{} : Bool}, decision: Shape.BoolCase<Nat.is_eq(position, previous)>,
99 earlier: @inside: {Nat.is_lt(position, previous) == True{} : Bool} ->
100 {ReduceLaws.word_at(Contracts.words(shape, fill, ops, previous), position) ==
101 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}) ->
102 {ReduceLaws.word_at(Contracts.words(shape, fill, ops, 1n+previous), position) ==
103 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}:
104 match decision:
105 case Shape.BFalse{equal}:
106 +before = Contracts.words(shape, fill, ops, previous)
107 +last = Replayed.value(shape, fill, ops, U32.from_nat(previous))
108 +inside = PointwiseProof.lt_succ_ne(position, previous, below, equal)
109 Equal.trans(Maybe<&2, F32>, ReduceLaws.word_at(List.append(&2, F32, before, [last]), position), ReduceLaws.word_at(before, position),
110 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))},
111 word_at_front(before, last, position,
112 %Equal.sym(Nat, List.length(&2, F32, before), previous, words_length(shape, fill, ops, previous)) :
113 {Nat.is_lt(position, _) == True{} : Bool}
114 inside),
115 earlier(inside))
116 case Shape.BTrue{equal}:
117 +before = Contracts.words(shape, fill, ops, previous)
118 +last = Replayed.value(shape, fill, ops, U32.from_nat(previous))
119 %Equal.sym(Nat, position, previous, Arith.eq_true(position, previous, Equal.sym(Bool, True{}, Nat.is_eq(position, previous), equal))) :
120 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{Replayed.value(shape, fill, ops, U32.from_nat(_))} : Maybe<&2, F32>}
121 %words_length(shape, fill, ops, previous) :
122 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{last} : Maybe<&2, F32>}
123 word_at_last(before, last)
124
125def words_at(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +position: Nat,
126 below: {Nat.is_lt(position, count) == True{} : Bool}) ->
127 {ReduceLaws.word_at(Contracts.words(shape, fill, ops, count), position) ==
128 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}:
129 match count:
130 case 0n:
131 Empty.absurd({ReduceLaws.word_at(Nil{}, position) == Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>},
132 Arith.lt_zero(position, below))
133 case 1n+previous:
134 words_step(shape, fill, ops, previous, position, below, Shape.bool_case(Nat.is_eq(position, previous)),
135 inside => words_at(shape, fill, ops, previous, position, inside))
136
137def matched_chosen(+same: Bool, +total: F32, +word: F32) ->
138 {ReduceLaws.matched(Some{total}, Some{word}, same) == Some{Fold.chosen(~Unit, ~Sum.add, Unit{}, same, total, word)} : Maybe<&2, F32>}:
139 match same:
140 case False{}:
141 {==}
142 case True{}:
143 {==}
144
145def fiber_total(+plan: Shape.BroadcastPlan, +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat,
146 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
147 {ReduceLaws.fiber(plan,
148 Contracts.words(Shape.right_shape(plan), fill, ops, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))), index, count) ==
149 Some{Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index},
150 Shape.right_shape(plan), fill, ops, count)} : Maybe<&2, F32>}:
151 match count:
152 case 0n:
153 {==}
154 case 1n+k:
155 +right = Shape.right_shape(plan)
156 +output = Shape.output_shape(plan)
157 +limit = Shape.logical_numel(Shape.dimensions(output))
158 +values = Contracts.words(right, fill, ops, Shape.logical_numel(Shape.dimensions(right)))
159 +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index}
160 +lm = ReduceFold.target_offset(plan, k)
161 +rm = ReduceFold.contribution_offset(plan, k)
162 +before = Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, k)
163 +word = Replayed.value(right, fill, ops, U32.from_nat(rm))
164 +lands = Nat.is_eq(lm, U32.to_nat(index))
165 +after : Maybe<&2, F32> = Some{Fold.chosen(~Unit, ~Sum.add, Unit{}, lands, before, word)}
166 +below = Conversion.le_successor_lt(k, limit, room)
167 +rm_below = BroadcastProof.right_walk_bound(plan, Shape.index_map(plan), k, below)
168 +lm_below = BroadcastProof.left_walk_bound(plan, Shape.index_map(plan), k, below)
169 %Equal.sym(Maybe<&1, U32 & U32>, Projection.offset_spec(plan, k), Some{(U32.from_nat(lm), U32.from_nat(rm))},
170 offsets_walk(plan, Shape.index_map(plan), k, BroadcastProof.output_inside(plan, k, below))) :
171 {ReduceLaws.contribute(values, index, ReduceLaws.fiber(plan, values, index, k), _) == after : Maybe<&2, F32>}
172 %Equal.sym(Maybe<&2, F32>, ReduceLaws.fiber(plan, values, index, k), Some{before},
173 fiber_total(plan, index, fill, ops, k, Conversion.le_predecessor(k, limit, room))) :
174 {ReduceLaws.matched(_, ReduceLaws.word_at(values, U32.to_nat(U32.from_nat(rm))), Nat.is_eq(U32.to_nat(U32.from_nat(lm)), U32.to_nat(index))) ==
175 after : Maybe<&2, F32>}
176 %Equal.sym(Nat, U32.to_nat(U32.from_nat(rm)), rm, PointwiseProof.word_exact(rm, right, rm_below)) :
177 {ReduceLaws.matched(Some{before}, ReduceLaws.word_at(values, _), Nat.is_eq(U32.to_nat(U32.from_nat(lm)), U32.to_nat(index))) ==
178 after : Maybe<&2, F32>}
179 %Equal.sym(Maybe<&2, F32>, ReduceLaws.word_at(values, rm), Some{word},
180 words_at(right, fill, ops, Shape.logical_numel(Shape.dimensions(right)), rm, rm_below)) :
181 {ReduceLaws.matched(Some{before}, _, Nat.is_eq(U32.to_nat(U32.from_nat(lm)), U32.to_nat(index))) == after : Maybe<&2, F32>}
182 %Equal.sym(Nat, U32.to_nat(U32.from_nat(lm)), lm, PointwiseProof.word_exact(lm, Shape.left_shape(plan), lm_below)) :
183 {ReduceLaws.matched(Some{before}, Some{word}, Nat.is_eq(_, U32.to_nat(index))) == after : Maybe<&2, F32>}
184 matched_chosen(lands, before, word)
185
186def source_inside(+plan: Shape.BroadcastPlan, +index: U32, +count: Nat,
187 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
188 {Pointwise.inside(~ReduceFold.Cell, ~ReduceFold.source, ReduceFold.Cell{plan, index}, count,
189 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}:
190 match count:
191 case 0n:
192 {==}
193 case 1n+previous:
194 +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index}
195 +limit = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
196 +bound = Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))
197 %Equal.sym(Bool, Pointwise.inside(~ReduceFold.Cell, ~ReduceFold.source, cell, previous, bound), True{},
198 source_inside(plan, index, previous, Conversion.le_predecessor(previous, limit, room))) :
199 {Bool.and(_, Nat.is_lt(ReduceFold.source(cell, previous), bound)) == True{} : Bool}
200 %Equal.sym(Bool, Nat.is_lt(ReduceFold.source(cell, previous), bound), True{},
201 BroadcastProof.right_walk_bound(plan, Shape.index_map(plan), previous, Conversion.le_successor_lt(previous, limit, room))) :
202 {Bool.and(True{}, _) == True{} : Bool}
203 {==}
204
205def contribution_room(+plan: Shape.BroadcastPlan,
206 same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>}) ->
207 {Nat.is_le(Shape.elements(Shape.right_shape(plan)), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}:
208 +right = Shape.right_shape(plan)
209 +limit = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
210 %Equal.sym(Nat, Shape.elements(right), Shape.numel(Shape.extents(right)), Shape.elements_exact(right)) :
211 {Nat.is_le(_, limit) == True{} : Bool}
212 %Equal.cong(List<&2, Nat>, Nat, extents => Shape.numel(extents), Shape.extents(Shape.output_shape(plan)), Shape.extents(right), same) :
213 {Nat.is_le(_, limit) == True{} : Bool}
214 Conversion.le_reflexive(limit)
215
216def fiber_contract(+plan: Shape.BroadcastPlan, +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>,
217 -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) ==
218 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat},
219 +expected: Maybe<&2, F32>, +total: F32, +count: Nat,
220 evidence: Records.FoldContract<Shape.right_shape(plan), total,
221 other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
222 Fold.run(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, count,
223 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>,
224 equation: {expected == Some{total} : Maybe<&2, F32>}) ->
225 Laws.FiberContract<Shape.right_shape(plan), expected,
226 other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
227 Fold.run(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, count,
228 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>:
229 Records.FoldContract{value, kept} = evidence
230 %Equal.sym(Maybe<&2, F32>, expected, Some{total}, equation) :
231 Laws.FiberContract<Shape.right_shape(plan), _,
232 other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
233 Fold.run(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, count,
234 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>
235 Laws.FiberContract{value, kept}
236
237def reduce_fold_fiber(+plan: Shape.BroadcastPlan,
238 same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>},
239 +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>,
240 -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) ==
241 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}) ->
242 Laws.FiberContract<Shape.right_shape(plan),
243 ReduceLaws.fiber(plan,
244 Contracts.words(Shape.right_shape(plan), fill, ops, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
245 index, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
246 other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
247 ReduceFold.cell(plan, index,
248 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>:
249 +right = Shape.right_shape(plan)
250 +count = Shape.elements(right)
251 +values = Contracts.words(right, fill, ops, Shape.logical_numel(Shape.dimensions(right)))
252 +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index}
253 +total = Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, count)
254 +room = contribution_room(plan, same)
255 fiber_contract(plan, index, fill, ops, fits,
256 ReduceLaws.fiber(plan, values, index, Shape.logical_numel(Shape.dimensions(right))), total, count,
257 FoldProof.contract(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, count,
258 source_inside(plan, index, count, room), fits),
259 %Shape.elements_exact(right) : {ReduceLaws.fiber(plan, values, index, _) == Some{total} : Maybe<&2, F32>}
260 fiber_total(plan, index, fill, ops, count, room))
261
262def Laws.reduce_fold_fiber(plan, same, index, fill, ops, fits):
263 reduce_fold_fiber(plan, same, index, fill, ops, fits)