src/autograd/tensor_backward_PROOF.bend
bendtorch-core@0.1.0.0 · 517 lines · 25 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
5import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
6import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
7import bendtorch-core@0.1.0.0/src/autograd/broadcast_pointwise_LAWS.bend as BroadcastPointwise
8import bendtorch-core@0.1.0.0/src/autograd/broadcast_pointwise_PROOF.bend as BroadcastPointwiseProof
9import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
10import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
11import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_LAWS.bend as Laws
12import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
13import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
14import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
15import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
16import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
17
18def eq_refl(+number: Nat) -> {True{} == Nat.is_eq(number, number) : Bool}:
19 match number:
20 case 0n:
21 {==}
22 case 1n+rest:
23 eq_refl(rest)
24
25def add_zero(+number: Nat) -> {number == Nat.add(number, 0n) : Nat}:
26 match number:
27 case 0n:
28 {==}
29 case 1n+rest:
30 Equal.cong(Nat, Nat, value => 1n+value, rest, Nat.add(rest, 0n), add_zero(rest))
31
32def add_succ(+left: Nat, -right: Nat) -> {1n+Nat.add(left, right) == Nat.add(left, 1n+right) : Nat}:
33 match left:
34 case 0n:
35 {==}
36 case 1n+rest:
37 Equal.cong(Nat, Nat, value => 1n+value, 1n+Nat.add(rest, right), Nat.add(rest, 1n+right), add_succ(rest, right))
38
39
40
41def older_ne(+older: Nat, -extra: Nat) -> {False{} == Nat.is_eq(Nat.add(older, 1n+extra), older) : Bool}:
42 match older:
43 case 0n:
44 {==}
45 case 1n+rest:
46 older_ne(rest, extra)
47
48def listed_back(entries: List<&1, Backward.Entry>) ->
49 {(entries, Laws.nodes_of(entries)) == Backward.listed(entries) : List<&1, Backward.Entry> & List<&2, Backward.Node>}:
50 match entries:
51 case Nil{}:
52 {==}
53 case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
54 %listed_back(rest) :
55 {(Backward.Entry{node, shape, value, saved, slot} <> rest, node <> Laws.nodes_of(rest)) ==
56 Backward.entry_listed(Backward.Entry{node, shape, value, saved, slot}, _) : List<&1, Backward.Entry> & List<&2, Backward.Node>}
57 {==}
58
59def Laws.negate_left_scalar(left, right):
60 {==}
61
62def Laws.tensor_vjp_leaf(saved, g):
63 {==}
64
65def Laws.tensor_vjp_constant(saved, g):
66 {==}
67
68def frozen_nt_unsent(+id: Nat, +a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Backward.Saved, g: Core.Dynamic,
69 computed: {Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))) == False{} : Bool}) ->
70 {Backward.nt_frozen(id, a, b, left, right, saved, g, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))) ==
71 Backward.unsent(id, Backward.vjp_matmul_nt(a, b, left, right, saved, g)) : Backward.Pass}:
72 %Equal.sym(Bool, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))), False{}, computed) :
73 {Backward.nt_frozen(id, a, b, left, right, saved, g, _) == Backward.unsent(id, Backward.vjp_matmul_nt(a, b, left, right, saved, g)) : Backward.Pass}
74 {==}
75
76def Laws.tensor_vjp_frozen_unsent(id, node, saved, g, computed):
77 match node:
78 case Backward.MatmulNT{+a, +b, +left, +right}:
79 frozen_nt_unsent(id, a, b, left, right, saved, g, computed)
80 case Backward.Leaf{}:
81 {==}
82 case Backward.Add{a, b, left, right}:
83 {==}
84 case Backward.Sub{a, b, left, right}:
85 {==}
86 case Backward.Mul{a, b, left, right}:
87 {==}
88 case Backward.Sum{a, input}:
89 {==}
90 case Backward.SumDims{a, input, kept}:
91 {==}
92 case Backward.Matmul{a, b, left, right}:
93 {==}
94 case Backward.Relu{a, input}:
95 {==}
96 case Backward.MseLoss{a, b, left, right}:
97 {==}
98 case Backward.CrossEntropy{a, input}:
99 {==}
100 case Backward.BatchedMatmul{a, b, left, right}:
101 {==}
102 case Backward.Slice{a, input, kept, plan}:
103 {==}
104 case Backward.Gathered{a, input, plan, entries}:
105 {==}
106 case Backward.Constant{}:
107 {==}
108 case Backward.Frozen{inner_id, inner}:
109 {==}
110
111
112def spared_one(+id: Nat, +target: Nat, +shape: Shape.Shape, -gradient: Core.Dynamic, -rest: List<&1, Backward.Contribution>,
113 ih: {Laws.reaches(id, Backward.kept_others(id, rest)) == False{} : Bool},
114 decision: Shape.BoolCase<Nat.is_eq(id, target)>) ->
115 {Laws.reaches(id, Backward.kept_one(target, shape, gradient, Nat.is_eq(id, target), Backward.kept_others(id, rest))) == False{} : Bool}:
116 match decision:
117 case Shape.BTrue{equal}:
118 %equal : {Laws.reaches(id, Backward.kept_one(target, shape, gradient, _, Backward.kept_others(id, rest))) == False{} : Bool}
119 ih
120 case Shape.BFalse{+equal}:
121 %equal : {Laws.reaches(id, Backward.kept_one(target, shape, gradient, _, Backward.kept_others(id, rest))) == False{} : Bool}
122 %equal : {Bool.or(_, Laws.reaches(id, Backward.kept_others(id, rest))) == False{} : Bool}
123 %Equal.sym(Bool, Laws.reaches(id, Backward.kept_others(id, rest)), False{}, ih) : {Bool.or(False{}, _) == False{} : Bool}
124 {==}
125
126def spared(+id: Nat, items: List<&1, Backward.Contribution>) -> {Laws.reaches(id, Backward.kept_others(id, items)) == False{} : Bool}:
127 match items:
128 case Nil{}:
129 {==}
130 case Con{Backward.Contribution{+target, +shape, gradient}, rest}:
131 spared_one(id, target, shape, gradient, rest, spared(id, rest), Shape.bool_case(Nat.is_eq(id, target)))
132
133def unsent_spared_result(+id: Nat, -saved: Backward.Saved, -g: Core.Dynamic,
134 contributions: Result<&2, &1, Backward.Error, List<&1, Backward.Contribution>>) ->
135 {Laws.sends_to(id, Backward.unsent(id, Backward.Pass{saved, g, contributions})) == False{} : Bool}:
136 match contributions:
137 case Fail{error}:
138 {==}
139 case Done{items}:
140 spared(id, items)
141
142def unsent_spared(+id: Nat, pass: Backward.Pass) -> {Laws.sends_to(id, Backward.unsent(id, pass)) == False{} : Bool}:
143 match pass:
144 case Backward.Pass{saved, g, contributions}:
145 unsent_spared_result(id, saved, g, contributions)
146
147def other_operand(x: Bool, y: Bool, both: {Bool.and(x, Bool.not(y)) == True{} : Bool}) -> {y == False{} : Bool}:
148 match x:
149 case False{}:
150 Empty.absurd({y == False{} : Bool}, Arith.clash(both))
151 case True{}:
152 match y:
153 case True{}:
154 Empty.absurd({True{} == False{} : Bool}, Arith.clash(both))
155 case False{}:
156 {==}
157
158def alone_spared(+id: Nat, +b: Nat, +right: Shape.Shape, -x: Core.Dynamic, -y: Core.Dynamic, -g: Core.Dynamic,
159 other: {Nat.is_eq(id, b) == False{} : Bool}, result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) ->
160 {Laws.sends_to(id, Backward.Pass{Backward.Operands{x, y}, g, Backward.alone(Backward.multiplied(b, right, result))}) == False{} : Bool}:
161 match result:
162 case Fail{error}:
163 {==}
164 case Done{triple}:
165 match triple:
166 case (kept_left, rest):
167 match rest:
168 case (kept_right, product):
169 %Equal.sym(Bool, Nat.is_eq(id, b), False{}, other) : {Bool.or(_, False{}) == False{} : Bool}
170 {==}
171
172def right_spared(+id: Nat, +b: Nat, +right: Shape.Shape, -y: Core.Dynamic, -g: Core.Dynamic,
173 other: {Nat.is_eq(id, b) == False{} : Bool}, copy: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) ->
174 {Laws.sends_to(id, Backward.nt_alone_second(y, g, Backward.nt_right(b, right, copy, pair))) == False{} : Bool}:
175 match pair:
176 case (left, flipped):
177 alone_spared(id, b, right, left, y, g, other, Backward.weight_product(copy, flipped))
178
179def first_spared(+id: Nat, +b: Nat, +right: Shape.Shape, x: Core.Dynamic, -y: Core.Dynamic,
180 other: {Nat.is_eq(id, b) == False{} : Bool}, pair: Core.Dynamic & Core.Dynamic) ->
181 {Laws.sends_to(id, Backward.nt_alone_first(b, right, x, y, pair)) == False{} : Bool}:
182 match pair:
183 case (g, copy):
184 right_spared(id, b, right, y, g, other, copy, Core.clone_dynamic(x))
185
186def right_alone_spared(+id: Nat, +b: Nat, +right: Shape.Shape, saved: Backward.Saved, g: Core.Dynamic,
187 other: {Nat.is_eq(id, b) == False{} : Bool}) ->
188 {Laws.sends_to(id, Backward.vjp_matmul_nt_right(b, right, saved, g)) == False{} : Bool}:
189 match saved:
190 case Backward.Unsaved{}:
191 {==}
192 case Backward.Operands{x, y}:
193 first_spared(id, b, right, x, y, other, Core.clone_dynamic(g))
194
195def frozen_nt_spared(+id: Nat, +a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Backward.Saved, g: Core.Dynamic,
196 decision: Shape.BoolCase<Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))>) ->
197 {Laws.sends_to(id, Backward.nt_frozen(id, a, b, left, right, saved, g, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))))) == False{} : Bool}:
198 match decision:
199 case Shape.BTrue{+equal}:
200 %equal : {Laws.sends_to(id, Backward.nt_frozen(id, a, b, left, right, saved, g, _)) == False{} : Bool}
201 right_alone_spared(id, b, right, saved, g,
202 other_operand(Nat.is_eq(id, a), Nat.is_eq(id, b), Equal.sym(Bool, True{}, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))), equal)))
203 case Shape.BFalse{equal}:
204 %equal : {Laws.sends_to(id, Backward.nt_frozen(id, a, b, left, right, saved, g, _)) == False{} : Bool}
205 unsent_spared(id, Backward.vjp_matmul_nt(a, b, left, right, saved, g))
206
207def Laws.tensor_vjp_frozen_spares(id, node, saved, g):
208 match node:
209 case Backward.MatmulNT{+a, +b, +left, +right}:
210 frozen_nt_spared(id, a, b, left, right, saved, g, Shape.bool_case(Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))))
211 case Backward.Leaf{}:
212 unsent_spared(id, Backward.vjp(Backward.Leaf{}, saved, g))
213 case Backward.Add{a, b, left, right}:
214 unsent_spared(id, Backward.vjp(Backward.Add{a, b, left, right}, saved, g))
215 case Backward.Sub{a, b, left, right}:
216 unsent_spared(id, Backward.vjp(Backward.Sub{a, b, left, right}, saved, g))
217 case Backward.Mul{a, b, left, right}:
218 unsent_spared(id, Backward.vjp(Backward.Mul{a, b, left, right}, saved, g))
219 case Backward.Sum{a, input}:
220 unsent_spared(id, Backward.vjp(Backward.Sum{a, input}, saved, g))
221 case Backward.SumDims{a, input, kept}:
222 unsent_spared(id, Backward.vjp(Backward.SumDims{a, input, kept}, saved, g))
223 case Backward.Matmul{a, b, left, right}:
224 unsent_spared(id, Backward.vjp(Backward.Matmul{a, b, left, right}, saved, g))
225 case Backward.Relu{a, input}:
226 unsent_spared(id, Backward.vjp(Backward.Relu{a, input}, saved, g))
227 case Backward.MseLoss{a, b, left, right}:
228 unsent_spared(id, Backward.vjp(Backward.MseLoss{a, b, left, right}, saved, g))
229 case Backward.CrossEntropy{a, input}:
230 unsent_spared(id, Backward.vjp(Backward.CrossEntropy{a, input}, saved, g))
231 case Backward.BatchedMatmul{a, b, left, right}:
232 unsent_spared(id, Backward.vjp(Backward.BatchedMatmul{a, b, left, right}, saved, g))
233 case Backward.Slice{a, input, kept, plan}:
234 unsent_spared(id, Backward.vjp(Backward.Slice{a, input, kept, plan}, saved, g))
235 case Backward.Gathered{a, input, plan, entries}:
236 unsent_spared(id, Backward.vjp(Backward.Gathered{a, input, plan, entries}, saved, g))
237 case Backward.Constant{}:
238 unsent_spared(id, Backward.vjp(Backward.Constant{}, saved, g))
239 case Backward.Frozen{inner_id, inner}:
240 unsent_spared(id, Backward.vjp(Backward.Frozen{inner_id, inner}, saved, g))
241
242def Laws.tensor_backward_unknown(entries, output, outside):
243 %listed_back(entries) :
244 {Backward.planned(output, _) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
245 Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
246 %Equal.sym(Bool, Nat.is_lt(output, List.length(&2, Backward.Node, Laws.nodes_of(entries))), False{}, outside) :
247 {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, _,
248 Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries)))) ==
249 (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) : Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
250 {==}
251
252def Laws.tensor_backward_dangling(entries, output, inside, unsound):
253 %listed_back(entries) :
254 {Backward.planned(output, _) == (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) :
255 Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
256 %Equal.sym(Bool, Nat.is_lt(output, List.length(&2, Backward.Node, Laws.nodes_of(entries))), True{}, inside) :
257 {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, _,
258 Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries)))) ==
259 (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) : Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
260 %Equal.sym(Bool, Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries))), False{}, unsound) :
261 {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, True{}, _) ==
262 (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) : Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
263 {==}
264
265def Laws.tensor_backward_success(entries, output, inside, sound):
266 %listed_back(entries) :
267 {Backward.planned(output, _) ==
268 Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output,
269 Backward.at_head(List.length(&2, Backward.Node, Laws.nodes_of(entries)), output))) :
270 Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
271 %Equal.sym(Bool, Nat.is_lt(output, List.length(&2, Backward.Node, Laws.nodes_of(entries))), True{}, inside) :
272 {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, _,
273 Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries)))) ==
274 Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output,
275 Backward.at_head(List.length(&2, Backward.Node, Laws.nodes_of(entries)), output))) :
276 Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
277 %Equal.sym(Bool, Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries))), True{}, sound) :
278 {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, True{}, _) ==
279 Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output,
280 Backward.at_head(List.length(&2, Backward.Node, Laws.nodes_of(entries)), output))) :
281 Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
282 {==}
283
284def and_true(+value: Bool) -> {Bool.and(value, True{}) == value : Bool}:
285 match value:
286 case True{}:
287 {==}
288 case False{}:
289 {==}
290
291def pair_bounded(+a: Nat, +b: Nat, +bound: Nat) ->
292 {Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound)) == Laws.bounded([a, b], bound) : Bool}:
293 Equal.cong(Bool, Bool, flag => Bool.and(Nat.is_lt(a, bound), flag), Nat.is_lt(b, bound), Bool.and(Nat.is_lt(b, bound), True{}),
294 Equal.sym(Bool, Bool.and(Nat.is_lt(b, bound), True{}), Nat.is_lt(b, bound), and_true(Nat.is_lt(b, bound))))
295
296def operands_bounded(+node: Backward.Node, +bound: Nat) ->
297 {Backward.operands_below(node, bound) == Laws.bounded(Laws.targets(node), bound) : Bool}:
298 match node:
299 case Backward.Leaf{}:
300 {==}
301 case Backward.Add{a, b, left, right}:
302 pair_bounded(a, b, bound)
303 case Backward.Sub{a, b, left, right}:
304 pair_bounded(a, b, bound)
305 case Backward.Mul{a, b, left, right}:
306 pair_bounded(a, b, bound)
307 case Backward.Sum{a, input}:
308 Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
309 case Backward.SumDims{a, input, kept}:
310 Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
311 case Backward.Matmul{a, b, left, right}:
312 pair_bounded(a, b, bound)
313 case Backward.MatmulNT{a, b, left, right}:
314 pair_bounded(a, b, bound)
315 case Backward.Relu{a, input}:
316 Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
317 case Backward.MseLoss{a, b, left, right}:
318 pair_bounded(a, b, bound)
319 case Backward.CrossEntropy{a, input}:
320 Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
321 case Backward.BatchedMatmul{a, b, left, right}:
322 pair_bounded(a, b, bound)
323 case Backward.Slice{a, input, kept, plan}:
324 Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
325 case Backward.Gathered{a, input, plan, entries}:
326 Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
327 case Backward.Constant{}:
328 {==}
329 case Backward.Frozen{id, inner}:
330 operands_bounded(inner, bound)
331
332def Laws.tensor_wired_matches_operand_bounds(graph):
333 match graph:
334 case Nil{}:
335 {==}
336 case Con{node, rest}:
337 %Equal.sym(Bool, Backward.operands_below(node, List.length(&2, Backward.Node, rest)),
338 Laws.bounded(Laws.targets(node), List.length(&2, Backward.Node, rest)), operands_bounded(node, List.length(&2, Backward.Node, rest))) :
339 {Bool.and(_, Backward.wired(rest, List.length(&2, Backward.Node, rest))) ==
340 Bool.and(Laws.bounded(Laws.targets(node), List.length(&2, Backward.Node, rest)), Laws.well_formed(rest)) : Bool}
341 Equal.cong(Bool, Bool, flag => Bool.and(Laws.bounded(Laws.targets(node), List.length(&2, Backward.Node, rest)), flag),
342 Backward.wired(rest, List.length(&2, Backward.Node, rest)), Laws.well_formed(rest), Laws.tensor_wired_matches_operand_bounds(rest))
343
344def Laws.tensor_seed_selected_output(newer, node, shape, value, saved, slot, rest, older, fits):
345 match newer:
346 case Nil{}:
347 %add_zero(older) :
348 {Backward.pass(Backward.Entry{node, shape, value, saved, slot} <> rest, 1n+_, older, Backward.at_head(1n+_, older)) ==
349 Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)) : Backward.Swept()}
350 %eq_refl(older) :
351 {Backward.pass(Backward.Entry{node, shape, value, saved, slot} <> rest, 1n+older, older, _) ==
352 Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)) : Backward.Swept()}
353 {==}
354 case Con{head, tail}:
355 %older_ne(older, List.length(&1, Backward.Entry, tail)) :
356 {Backward.pass(head <> List.append(&1, Backward.Entry, tail, Backward.Entry{node, shape, value, saved, slot} <> rest),
357 1n+Nat.add(older, 1n+List.length(&1, Backward.Entry, tail)), older, _) ==
358 Backward.finished(head, Gradient.Unset{},
359 Laws.newer_unset(tail, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)))) : Backward.Swept()}
360 %add_succ(older, List.length(&1, Backward.Entry, tail)) :
361 {Backward.pass(head <> List.append(&1, Backward.Entry, tail, Backward.Entry{node, shape, value, saved, slot} <> rest),
362 1n+_, older, False{}) ==
363 Backward.finished(head, Gradient.Unset{},
364 Laws.newer_unset(tail, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)))) : Backward.Swept()}
365 Equal.cong(Backward.Swept(), Backward.Swept(), swept => Backward.finished(head, Gradient.Unset{}, swept),
366 Backward.pass(List.append(&1, Backward.Entry, tail, Backward.Entry{node, shape, value, saved, slot} <> rest),
367 1n+Nat.add(older, List.length(&1, Backward.Entry, tail)), older,
368 Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, tail)), older)),
369 Laws.newer_unset(tail, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0))),
370 Laws.tensor_seed_selected_output(tail, node, shape, value, saved, slot, rest, older, fits))
371
372def Laws.tensor_seeded_sweep(node, shape, value, saved, slot, rest, below, ones):
373 {==}
374
375def Laws.tensor_unreached_sends_nothing(node, shape, value, saved, slot, rest, pending, below):
376 {==}
377
378def Laws.tensor_deliver_in_order(target, shape, gradient, rest, count, slots):
379 {==}
380
381def Laws.tensor_bump_target(newer, slot, rest, shape, gradient, older, fits):
382 match newer:
383 case Nil{}:
384 %add_zero(older) :
385 {Backward.bump(slot <> rest, 1n+_, older, shape, gradient, Backward.at_head(1n+_, older)) ==
386 Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)) : Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
387 %eq_refl(older) :
388 {Backward.bump(slot <> rest, 1n+older, older, shape, gradient, _) ==
389 Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)) : Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
390 {==}
391 case Con{head, tail}:
392 %older_ne(older, List.length(&1, Gradient.Slot, tail)) :
393 {Backward.bump(head <> List.append(&1, Gradient.Slot, tail, slot <> rest),
394 1n+Nat.add(older, 1n+List.length(&1, Gradient.Slot, tail)), older, shape, gradient, _) ==
395 Backward.prepend(Gradient.Slot, head,
396 Laws.relisted(tail, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)))) :
397 Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
398 %add_succ(older, List.length(&1, Gradient.Slot, tail)) :
399 {Backward.bump(head <> List.append(&1, Gradient.Slot, tail, slot <> rest), 1n+_, older, shape, gradient, False{}) ==
400 Backward.prepend(Gradient.Slot, head,
401 Laws.relisted(tail, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)))) :
402 Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
403 Equal.cong(Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>, Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>,
404 result => Backward.prepend(Gradient.Slot, head, result),
405 Backward.bump(List.append(&1, Gradient.Slot, tail, slot <> rest), 1n+Nat.add(older, List.length(&1, Gradient.Slot, tail)),
406 older, shape, gradient, Backward.at_head(1n+Nat.add(older, List.length(&1, Gradient.Slot, tail)), older)),
407 Laws.relisted(tail, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient))),
408 Laws.tensor_bump_target(tail, slot, rest, shape, gradient, older, fits))
409
410def Laws.tensor_merge_accumulates(node, shape, value, saved, slot, rest, g, more):
411 {==}
412
413def Laws.tensor_merge_keeps_unreached(node, shape, value, saved, slot, rest, more):
414 {==}
415
416def Laws.tensor_zero_grad_empties(entries):
417 match entries:
418 case Nil{}:
419 {==}
420 case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
421 Equal.cong(List<&1, Gradient.Slot>, List<&1, Gradient.Slot>, slots => Gradient.Unset{} <> slots,
422 Laws.tape_slots(Backward.zero_grad(Backward.Tape{rest})), Backward.unsets(List.length(&1, Backward.Entry, rest)),
423 Laws.tensor_zero_grad_empties(rest))
424
425def Laws.tensor_zero_grad_keeps_nodes(entries):
426 match entries:
427 case Nil{}:
428 {==}
429 case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
430 Equal.cong(List<&2, Backward.Node>, List<&2, Backward.Node>, nodes => node <> nodes,
431 Laws.tape_nodes(Backward.zero_grad(Backward.Tape{rest})), Laws.nodes_of(rest),
432 Laws.tensor_zero_grad_keeps_nodes(rest))
433
434def truth(value: Bool) -> Data:
435 match value:
436 case True{}:
437 Unit
438 case False{}:
439 Empty
440
441def false_true(equal: {False{} == True{} : Bool}) -> Empty:
442 %Equal.sym(Bool, False{}, True{}, equal) : truth(_)
443 Unit{}
444
445def sum_result(+a: Nat, +input: Shape.Shape, -shape: Shape.Shape, -cell: Nat -> F32, -word: U32 -> F32, -owner: Core.Dynamic,
446 result: Result<&2, &1, Broadcast.Error, Core.Dynamic>,
447 contract: BroadcastPointwise.BroadcastToContract<input, shape, cell, word, (owner, result)>) ->
448 Laws.SoleContract<a, input, cell, Backward.expanded(a, input, (owner, result))>:
449 match result:
450 case Fail{error}:
451 BroadcastPointwise.BroadcastToContract{successful, dimensions, value, kept} = contract
452 Empty.absurd(Laws.SoleContract<a, input, cell, Backward.expanded(a, input, (owner, Fail{error}))>, false_true(successful))
453 case Done{Core.Packed{output, tensor}}:
454 BroadcastPointwise.BroadcastToContract{successful, dimensions, value, kept} = contract
455 Laws.SoleContract{{==}, dimensions, value}
456
457
458
459def sum_contract(+a: Nat, +input: Shape.Shape, -shape: Shape.Shape, -cell: Nat -> F32, -word: U32 -> F32,
460 pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>,
461 contract: BroadcastPointwise.BroadcastToContract<input, shape, cell, word, pair>) ->
462 Laws.SoleContract<a, input, cell, Backward.expanded(a, input, pair)>:
463 (owner, result) = pair
464 sum_result(a, input, shape, cell, word, owner, result, contract)
465
466def Laws.tensor_sum_vjp_pointwise(a, shape, input, plan, planned, same, saved, fill, ops, fits):
467 sum_contract(a, input, shape,
468 index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
469 other => Replayed.value(shape, fill, ops, other),
470 Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, input),
471 BroadcastPointwise.broadcast_to_pointwise(shape, input, plan, planned, same, fill, ops, fits))
472
473
474def restored_contract(-a: Nat, -target: Shape.Shape, -cell: Nat -> F32, +shape: Shape.Shape, pass: Backward.Pass,
475 contract: Laws.SoleContract<a, target, cell, pass>) -> Laws.SoleContract<a, target, cell, Backward.restored(shape, pass)>:
476 match pass:
477 case Backward.Pass{saved, g, Fail{error}}:
478 Laws.SoleContract{target_id, dimensions, value} = contract
479 Laws.SoleContract{target_id, dimensions, value}
480 case Backward.Pass{saved, g, Done{Nil{}}}:
481 Laws.SoleContract{target_id, dimensions, value} = contract
482 Laws.SoleContract{target_id, dimensions, value}
483 case Backward.Pass{saved, g, Done{Con{contribution, Nil{}}}}:
484 Laws.SoleContract{target_id, dimensions, value} = contract
485 Laws.SoleContract{target_id, dimensions, value}
486 case Backward.Pass{saved, g, Done{Con{contribution, Con{other, rest}}}}:
487 Laws.SoleContract{target_id, dimensions, value} = contract
488 Laws.SoleContract{target_id, dimensions, value}
489
490def sum_dims_decided(+a: Nat, +shape: Shape.Shape, +input: Shape.Shape, +kept: Shape.Shape, +plan: Shape.BroadcastPlan,
491 planned: {Shape.broadcast_checked(kept, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
492 same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>},
493 counted: {Nat.is_eq(Shape.elements(shape), Shape.elements(kept)) == True{} : Bool},
494 saved: Backward.Saved, +fill: F32, +ops: List<&2, StorageLaws.Op>,
495 -fits: {Lifecycle.logical_length(Replayed.storage(kept, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
496 decision: Shape.BoolCase<Nat.is_eq(Shape.elements(shape), Shape.elements(kept))>) ->
497 Laws.SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
498 Backward.unsqueezed(a, input, kept, shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}, decision)>:
499 match decision:
500 case Shape.BFalse{equal}:
501 Empty.absurd(Laws.SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
502 Backward.unsqueezed(a, input, kept, shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}, Shape.BFalse{equal})>,
503 Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(Shape.elements(shape), Shape.elements(kept)), True{}, equal, counted)))
504 case Shape.BTrue{equal}:
505 restored_contract(a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))), shape,
506 Backward.vjp_sum(a, input, Transforms.reshape_equal(shape, kept, Core.TTensor{Replayed.storage(kept, fill, ops), fits},
507 Transforms.elements_equal(shape, kept, equal))),
508 Laws.tensor_sum_vjp_pointwise(a, kept, input, plan, planned, same, saved, fill, ops,
509 Equal.trans(Nat, Lifecycle.logical_length(Replayed.storage(kept, fill, ops)), Shape.logical_numel(Shape.dimensions(shape)),
510 Shape.logical_numel(Shape.dimensions(kept)), fits, Transforms.elements_equal(shape, kept, equal))))
511
512def Laws.tensor_sum_dims_vjp_pointwise(a, shape, input, kept, plan, planned, same, counted, saved, fill, ops, fits):
513 sum_dims_decided(a, shape, input, kept, plan, planned, same, counted, saved, fill, ops, fits,
514 Shape.bool_case(Nat.is_eq(Shape.elements(shape), Shape.elements(kept))))
515
516def Laws.tensor_sum_dims_wraps(a, axes, keepdim, entries, input, shape):
517 {==}