src/autograd/backward_LAWS.bend

bendtorch-core@0.1.0.0 · 456 lines · 61 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/autograd/backward.bend as Backward
3
4# Reachability, stated on flags: a node is reached when it is the selected
5# output, or an operand of a reached node newer than it. Nodes newer than the
6# output are never reached.
7
8def marked(flag: Bool, +rest: List<&2, Bool>, rec: Unit -> List<&2, Bool>, same: Bool) -> List<&2, Bool>:
9  match same:
10    case True{}:
11      True{} <> rest
12    case False{}:
13      flag <> rec(Unit{})
14
15def mark(flags: List<&2, Bool>, +count: Nat, +target: Nat) -> List<&2, Bool>:
16  match flags count:
17    case Nil{} count:
18      Nil{}
19    case Con{flag, rest} 0n:
20      flag <> rest
21    case Con{flag, +rest} 1n+below:
22      marked(flag, rest, unit => mark(rest, below, target), Nat.is_eq(below, target))
23
24def targets(node: Backward.Node) -> List<&2, Nat>:
25  match node:
26    case Backward.Leaf{}:
27      []
28    case Backward.Add{a, b}:
29      [a, b]
30    case Backward.Sub{a, b}:
31      [a, b]
32    case Backward.Mul{a, b, left, right}:
33      [a, b]
34    case Backward.Neg{a}:
35      [a]
36
37def mark_all(ids: List<&2, Nat>, +count: Nat, flags: List<&2, Bool>) -> List<&2, Bool>:
38  match ids:
39    case Nil{}:
40      flags
41    case Con{+id, rest}:
42      mark_all(rest, count, mark(flags, count, id))
43
44def spread(node: Backward.Node, flag: Bool, +count: Nat, flags: List<&2, Bool>) -> List<&2, Bool>:
45  match flag:
46    case False{}:
47      flags
48    case True{}:
49      mark_all(targets(node), count, flags)
50
51def reach(nodes: List<&2, Backward.Node>, flags: List<&2, Bool>, +count: Nat) -> List<&2, Bool>:
52  match nodes flags count:
53    case Nil{} flags count:
54      flags
55    case Con{node, rest} Nil{} count:
56      Nil{}
57    case Con{node, rest} Con{flag, pending} 0n:
58      Nil{}
59    case Con{node, rest} Con{+flag, pending} 1n+below:
60      flag <> reach(rest, spread(node, flag, below, pending), below)
61
62def falses(count: Nat) -> List<&2, Bool>:
63  match count:
64    case 0n:
65      Nil{}
66    case 1n+rest:
67      False{} <> falses(rest)
68
69def reached_chosen(node: Backward.Node, +rest: List<&2, Backward.Node>, +below: Nat, rec: Unit -> List<&2, Bool>, selected: Bool) ->
70  List<&2, Bool>:
71  match selected:
72    case True{}:
73      reach(node <> rest, True{} <> falses(below), 1n+below)
74    case False{}:
75      False{} <> rec(Unit{})
76
77def reached(nodes: List<&2, Backward.Node>, +count: Nat, +output: Nat) -> List<&2, Bool>:
78  match nodes count:
79    case Nil{} count:
80      Nil{}
81    case Con{node, rest} 0n:
82      Nil{}
83    case Con{node, +rest} 1n+below:
84      reached_chosen(node, rest, below, unit => reached(rest, below, output), Nat.is_eq(below, output))
85
86# Operand bounds, stated without the engine's validator: a node may name only
87# nodes older than itself, so each operand id must be below the number of
88# older nodes. An operand that is the node itself, a newer node or a missing
89# node fails this bound.
90def bounded(ids: List<&2, Nat>, +bound: Nat) -> Bool:
91  match ids:
92    case Nil{}:
93      True{}
94    case Con{+id, rest}:
95      Bool.and(Nat.is_lt(id, bound), bounded(rest, bound))
96
97def well_formed(nodes: List<&2, Backward.Node>) -> Bool:
98  match nodes:
99    case Nil{}:
100      True{}
101    case Con{node, +rest}:
102      Bool.and(bounded(targets(node), List.length(&2, Backward.Node, rest)), well_formed(rest))
103
104def entry_node(entry: Backward.Entry) -> Backward.Node:
105  Backward.Entry{node, value, slot} = entry
106  node
107
108def somes(gradients: List<&2, Maybe<&2, F32>>) -> List<&2, Bool>:
109  match gradients:
110    case Nil{}:
111      Nil{}
112    case Con{gradient, rest}:
113      Maybe.is_some(&2, F32, gradient) <> somes(rest)
114
115def nodes_of(tape: Backward.Tape) -> List<&2, Backward.Node>:
116  Backward.Tape{entries} = tape
117  Backward.nodes(entries)
118
119def slots(entries: List<&2, Backward.Entry>) -> List<&2, Maybe<&2, F32>>:
120  match entries:
121    case Nil{}:
122      Nil{}
123    case Con{Backward.Entry{node, value, slot}, rest}:
124      slot <> slots(rest)
125
126def tape_slots(tape: Backward.Tape) -> List<&2, Maybe<&2, F32>>:
127  Backward.Tape{entries} = tape
128  slots(entries)
129
130def slots_after(result: Result<&2, &2, Backward.Error, Backward.Tape>) -> Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>:
131  match result:
132    case Fail{error}:
133      Fail{error}
134    case Done{tape}:
135      Done{tape_slots(tape)}
136
137def grad_after(result: Result<&2, &2, Backward.Error, Backward.Tape>, +id: Nat) -> Result<&2, &2, Backward.Error, Maybe<&2, F32>>:
138  match result:
139    case Fail{error}:
140      Fail{error}
141    case Done{tape}:
142      Backward.grad(tape, id)
143
144def again(result: Result<&2, &2, Backward.Error, Backward.Tape>, +output: Nat) -> Result<&2, &2, Backward.Error, Backward.Tape>:
145  match result:
146    case Fail{error}:
147      Fail{error}
148    case Done{tape}:
149      Backward.backward(tape, output)
150
151def observe(+tape: Backward.Tape, +output: Nat) ->
152  Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>:
153  (Backward.value(tape, output), Backward.gradients(tape, output))
154
155def tape_of(pair: Backward.Tape & Nat) -> Backward.Tape:
156  (tape, id) = pair
157  tape
158
159def recorded(pair: Backward.Tape & Result<&2, &2, Backward.Error, Nat>) -> Backward.Tape:
160  (tape, result) = pair
161  tape
162
163def program_binary(kind: Backward.Binary, x: F32, y: F32) -> Backward.Tape:
164  recorded(Backward.binary(tape_of(Backward.leaf(tape_of(Backward.leaf(Backward.empty(), x)), y)), kind, 0n, 1n))
165
166def program_neg(x: F32) -> Backward.Tape:
167  recorded(Backward.neg(tape_of(Backward.leaf(Backward.empty(), x)), 0n))
168
169# x = 0, y = 1, z = 2, x * y = 3, z + z = 4. Node 4 is newer than the output
170# 3 and uses z twice; z is not used by 3.
171def program_unused(x: F32, y: F32, z: F32) -> Backward.Tape:
172  recorded(Backward.add(recorded(Backward.mul(
173    tape_of(Backward.leaf(tape_of(Backward.leaf(tape_of(Backward.leaf(Backward.empty(), x)), y)), z)), 0n, 1n)), 2n, 2n))
174
175# x = 0, a = 1, b = 2, c = 3, p = x * a = 4, q = x * b = 5, r = x * c = 6,
176# s = p + q = 7, t = s + r = 8. x has three consumers.
177def program_fan_out(x: F32, a: F32, b: F32, c: F32) -> Backward.Tape:
178  recorded(Backward.add(recorded(Backward.add(recorded(Backward.mul(recorded(Backward.mul(recorded(Backward.mul(
179    tape_of(Backward.leaf(tape_of(Backward.leaf(tape_of(Backward.leaf(tape_of(Backward.leaf(Backward.empty(), x)), a)), b)), c)),
180    0n, 1n)), 0n, 2n)), 0n, 3n)), 4n, 5n)), 7n, 6n))
181
182# Vector-Jacobian products: the contribution each operand receives when the
183# node's gradient is g, in operand order.
184
185law vjp_leaf:
186  for +g: F32
187  {Backward.vjp(Backward.Leaf{}, g) == [] : List<&2, Backward.Contribution>}
188
189law vjp_add:
190  for +a: Nat
191  for +b: Nat
192  for +g: F32
193  {Backward.vjp(Backward.Add{a, b}, g) == [Backward.Contribution{a, g}, Backward.Contribution{b, g}] : List<&2, Backward.Contribution>}
194
195law vjp_sub:
196  for +a: Nat
197  for +b: Nat
198  for +g: F32
199  {Backward.vjp(Backward.Sub{a, b}, g) == [Backward.Contribution{a, g}, Backward.Contribution{b, F32.neg(g)}] : List<&2, Backward.Contribution>}
200
201# d(a*b)/da = b and d(a*b)/db = a, with the operand values saved at record
202# time: the first operand receives g * right, the second g * left.
203law vjp_mul:
204  for +a: Nat
205  for +b: Nat
206  for +left: F32
207  for +right: F32
208  for +g: F32
209  {Backward.vjp(Backward.Mul{a, b, left, right}, g) ==
210    [Backward.Contribution{a, F32.mul(g, right)}, Backward.Contribution{b, F32.mul(g, left)}] : List<&2, Backward.Contribution>}
211
212law vjp_neg:
213  for +a: Nat
214  for +g: F32
215  {Backward.vjp(Backward.Neg{a}, g) == [Backward.Contribution{a, F32.neg(g)}] : List<&2, Backward.Contribution>}
216
217# Recording through the public API and running backward gives each op's
218# forward value and its VJP at the leaves (gradients newest first).
219
220law backward_add:
221  for +x: F32
222  for +y: F32
223  {observe(program_binary(Backward.Plus{}, x, y), 2n) ==
224    (Done{F32.add(x, y)}, Done{[Some{1.0}, Some{1.0}, Some{1.0}]}) :
225    Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
226
227law backward_sub:
228  for +x: F32
229  for +y: F32
230  {observe(program_binary(Backward.Minus{}, x, y), 2n) ==
231    (Done{F32.sub(x, y)}, Done{[Some{1.0}, Some{F32.neg(1.0)}, Some{1.0}]}) :
232    Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
233
234law backward_mul:
235  for +x: F32
236  for +y: F32
237  {observe(program_binary(Backward.Times{}, x, y), 2n) ==
238    (Done{F32.mul(x, y)}, Done{[Some{1.0}, Some{F32.mul(1.0, x)}, Some{F32.mul(1.0, y)}]}) :
239    Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
240
241law backward_neg:
242  for +x: F32
243  {observe(program_neg(x), 1n) ==
244    (Done{F32.neg(x)}, Done{[Some{1.0}, Some{F32.neg(1.0)}]}) :
245    Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
246
247# The seed is the selected output (C1). For any graph, the nodes newer than
248# the output get None, the output gets exactly 1.0, and the older nodes get
249# what the sweep sends from the output's VJP, whatever the newer nodes are.
250law seed_selected_output:
251  for +newer: List<&2, Backward.Node>
252  for +node: Backward.Node
253  for +rest: List<&2, Backward.Node>
254  {Backward.pass(List.append(&2, Backward.Node, newer, node <> rest),
255      List.length(&2, Backward.Node, List.append(&2, Backward.Node, newer, node <> rest)), List.length(&2, Backward.Node, rest)) ==
256    List.append(&2, Maybe<&2, F32>, Backward.nones(List.length(&2, Backward.Node, newer)),
257      Some{1.0} <> Backward.sweep(rest,
258        Backward.contribute(node, Some{1.0}, List.length(&2, Backward.Node, rest), Backward.nones(List.length(&2, Backward.Node, rest))),
259        List.length(&2, Backward.Node, rest))) : List<&2, Maybe<&2, F32>>}
260
261# Only the reachable graph: a node has a gradient exactly when the selected
262# output reaches it. Every other node, newer or unused, reads None.
263law reached_exactly:
264  for +graph: List<&2, Backward.Node>
265  for +output: Nat
266  {somes(Backward.pass(graph, List.length(&2, Backward.Node, graph), output)) ==
267    reached(graph, List.length(&2, Backward.Node, graph), output) : List<&2, Bool>}
268
269# With x, y and z recorded and z only used by a node newer than the output
270# x * y, backward of x * y seeds the output, leaves z and the newer node None,
271# and gives x and y their VJP.
272law unused_leaf_none:
273  for +x: F32
274  for +y: F32
275  for +z: F32
276  {slots_after(Backward.backward(program_unused(x, y, z), 3n)) ==
277    Done{[None{}, Some{1.0}, None{}, Some{F32.mul(1.0, x)}, Some{F32.mul(1.0, y)}]} :
278    Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
279
280# Fan-out: a node's first contribution moves in, and each later one is added
281# to the held total with one F32 add, total first.
282law accumulate_moves_in:
283  for contribution: Maybe<&2, F32>
284  {Backward.accumulate(None{}, contribution) == contribution : Maybe<&2, F32>}
285
286law accumulate_adds_total_first:
287  for +held: F32
288  for +value: F32
289  {Backward.accumulate(Some{held}, Some{value}) == Some{F32.add(held, value)} : Maybe<&2, F32>}
290
291law accumulate_keeps_without_contribution:
292  for +held: F32
293  {Backward.accumulate(Some{held}, None{}) == Some{held} : Maybe<&2, F32>}
294
295# A contribution to node target reaches exactly the slot of that node and no
296# other, for any list of slots.
297law bump_target:
298  for +newer: List<&2, Maybe<&2, F32>>
299  for +slot: Maybe<&2, F32>
300  for +rest: List<&2, Maybe<&2, F32>>
301  for +value: F32
302  {Backward.bump(List.append(&2, Maybe<&2, F32>, newer, slot <> rest),
303      List.length(&2, Maybe<&2, F32>, List.append(&2, Maybe<&2, F32>, newer, slot <> rest)), List.length(&2, Maybe<&2, F32>, rest), value) ==
304    List.append(&2, Maybe<&2, F32>, newer, Backward.accumulate(slot, Some{value}) <> rest) : List<&2, Maybe<&2, F32>>}
305
306# The contributions of x's three consumers arrive newest consumer first, and
307# each is added to the held total.
308law fan_out_arrival_order:
309  for +x: F32
310  for +a: F32
311  for +b: F32
312  for +c: F32
313  {grad_after(Backward.backward(program_fan_out(x, a, b, c), 8n), 0n) ==
314    Done{Some{F32.add(F32.add(F32.mul(1.0, c), F32.mul(1.0, b)), F32.mul(1.0, a))}} :
315    Result<&2, &2, Backward.Error, Maybe<&2, F32>>}
316
317# Ids (C2): id k reads exactly the element with k older elements, and an id
318# the list does not hold reads None, never another element.
319law lookup_exact:
320  for -A: Data
321  for +newer: List<&2, A>
322  for +item: A
323  for +rest: List<&2, A>
324  {Backward.lookup(A, List.append(&2, A, newer, item <> rest), List.length(&2, A, rest)) == Some{item} : Maybe<&2, A>}
325
326law lookup_unknown:
327  for -A: Data
328  for +items: List<&2, A>
329  for +id: Nat
330  for outside: {Nat.is_lt(id, List.length(&2, A, items)) == False{} : Bool}
331  {Backward.lookup(A, items, id) == None{} : Maybe<&2, A>}
332
333law bump_unknown:
334  for +gradients: List<&2, Maybe<&2, F32>>
335  for +target: Nat
336  for +value: F32
337  for +outside: {Nat.is_lt(target, List.length(&2, Maybe<&2, F32>, gradients)) == False{} : Bool}
338  {Backward.bump(gradients, List.length(&2, Maybe<&2, F32>, gradients), target, value) == gradients : List<&2, Maybe<&2, F32>>}
339
340law value_unknown:
341  for +entries: List<&2, Backward.Entry>
342  for +id: Nat
343  for outside: {Nat.is_lt(id, List.length(&2, Backward.Entry, entries)) == False{} : Bool}
344  {Backward.value(Backward.Tape{entries}, id) == Fail{Backward.UnknownNode{}} : Result<&2, &2, Backward.Error, F32>}
345
346law grad_unknown:
347  for +entries: List<&2, Backward.Entry>
348  for +id: Nat
349  for outside: {Nat.is_lt(id, List.length(&2, Backward.Entry, entries)) == False{} : Bool}
350  {Backward.grad(Backward.Tape{entries}, id) == Fail{Backward.UnknownNode{}} : Result<&2, &2, Backward.Error, Maybe<&2, F32>>}
351
352# Recording with an operand id the tape does not hold is refused with
353# UnknownNode, and the tape comes back unchanged.
354law record_unknown_left:
355  for +entries: List<&2, Backward.Entry>
356  for kind: Backward.Binary
357  for +a: Nat
358  for +b: Nat
359  for outside: {Nat.is_lt(a, List.length(&2, Backward.Entry, entries)) == False{} : Bool}
360  {Backward.binary(Backward.Tape{entries}, kind, a, b) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
361    Backward.Tape & Result<&2, &2, Backward.Error, Nat>}
362
363law record_unknown_right:
364  for +entries: List<&2, Backward.Entry>
365  for kind: Backward.Binary
366  for +a: Nat
367  for +b: Nat
368  for outside: {Nat.is_lt(b, List.length(&2, Backward.Entry, entries)) == False{} : Bool}
369  {Backward.binary(Backward.Tape{entries}, kind, a, b) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
370    Backward.Tape & Result<&2, &2, Backward.Error, Nat>}
371
372law neg_unknown:
373  for +entries: List<&2, Backward.Entry>
374  for +a: Nat
375  for outside: {Nat.is_lt(a, List.length(&2, Backward.Entry, entries)) == False{} : Bool}
376  {Backward.neg(Backward.Tape{entries}, a) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
377    Backward.Tape & Result<&2, &2, Backward.Error, Nat>}
378
379# backward of an output the tape does not hold is refused with UnknownNode.
380law backward_unknown:
381  for +entries: List<&2, Backward.Entry>
382  for +output: Nat
383  for outside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == False{} : Bool}
384  {Backward.backward(Backward.Tape{entries}, output) == Fail{Backward.UnknownNode{}} : Result<&2, &2, Backward.Error, Backward.Tape>}
385
386# A tape with a node that refers to itself, to a newer node or to a missing
387# node is refused with DanglingOperand.
388law backward_dangling:
389  for +entries: List<&2, Backward.Entry>
390  for +output: Nat
391  for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == True{} : Bool}
392  for unsound: {Backward.wired(Backward.nodes(entries), List.length(&2, Backward.Entry, entries)) == False{} : Bool}
393  {Backward.backward(Backward.Tape{entries}, output) == Fail{Backward.DanglingOperand{}} : Result<&2, &2, Backward.Error, Backward.Tape>}
394
395# The engine's wiring check is exactly the operand-bound specification above.
396law wired_matches_operand_bounds:
397  for +graph: List<&2, Backward.Node>
398  {Backward.wired(graph, List.length(&2, Backward.Node, graph)) == well_formed(graph) : Bool}
399
400# An entry with an operand that is itself, a newer entry or a missing entry
401# makes backward of any output on the tape fail with DanglingOperand.
402law backward_rejects_unbound_operand:
403  for +newer: List<&2, Backward.Entry>
404  for +entry: Backward.Entry
405  for +rest: List<&2, Backward.Entry>
406  for +output: Nat
407  for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, List.append(&2, Backward.Entry, newer, entry <> rest))) == True{} : Bool}
408  for unbound: {bounded(targets(entry_node(entry)), List.length(&2, Backward.Entry, rest)) == False{} : Bool}
409  {Backward.backward(Backward.Tape{List.append(&2, Backward.Entry, newer, entry <> rest)}, output) == Fail{Backward.DanglingOperand{}} :
410    Result<&2, &2, Backward.Error, Backward.Tape>}
411
412# A tape whose every operand is bound is accepted for every output it holds.
413law backward_accepts_bound_operands:
414  for +entries: List<&2, Backward.Entry>
415  for +output: Nat
416  for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == True{} : Bool}
417  for bound: {well_formed(Backward.nodes(entries)) == True{} : Bool}
418  {Backward.backward(Backward.Tape{entries}, output) ==
419    Done{Backward.Tape{Backward.merge(entries, Backward.pass(Backward.nodes(entries), List.length(&2, Backward.Entry, entries), output))}} :
420    Result<&2, &2, Backward.Error, Backward.Tape>}
421
422# Otherwise backward accumulates the gradients of pass into the slots.
423law backward_success:
424  for +entries: List<&2, Backward.Entry>
425  for +output: Nat
426  for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == True{} : Bool}
427  for sound: {Backward.wired(Backward.nodes(entries), List.length(&2, Backward.Entry, entries)) == True{} : Bool}
428  {Backward.backward(Backward.Tape{entries}, output) ==
429    Done{Backward.Tape{Backward.merge(entries, Backward.pass(Backward.nodes(entries), List.length(&2, Backward.Entry, entries), output))}} :
430    Result<&2, &2, Backward.Error, Backward.Tape>}
431
432# Repeated backward, like PyTorch's backward(retain_graph=True): the second
433# backward computes the same gradients and accumulates them again.
434law backward_twice_accumulates:
435  for +entries: List<&2, Backward.Entry>
436  for +output: Nat
437  for +found: List<&2, Maybe<&2, F32>>
438  for +computed: {Backward.gradients(Backward.Tape{entries}, output) == Done{found} : Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
439  {again(Backward.backward(Backward.Tape{entries}, output), output) ==
440    Done{Backward.Tape{Backward.merge(Backward.merge(entries, found), found)}} : Result<&2, &2, Backward.Error, Backward.Tape>}
441
442law backward_twice_mul:
443  for +x: F32
444  for +y: F32
445  {slots_after(again(Backward.backward(program_binary(Backward.Times{}, x, y), 2n), 2n)) ==
446    Done{[Some{F32.add(1.0, 1.0)}, Some{F32.add(F32.mul(1.0, x), F32.mul(1.0, x))}, Some{F32.add(F32.mul(1.0, y), F32.mul(1.0, y))}]} :
447    Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>}
448
449# zero_grad empties every slot and keeps every node.
450law zero_grad_empties:
451  for +entries: List<&2, Backward.Entry>
452  {tape_slots(Backward.zero_grad(Backward.Tape{entries})) == Backward.nones(List.length(&2, Backward.Entry, entries)) : List<&2, Maybe<&2, F32>>}
453
454law zero_grad_keeps_nodes:
455  for +entries: List<&2, Backward.Entry>
456  {nodes_of(Backward.zero_grad(Backward.Tape{entries})) == Backward.nodes(entries) : List<&2, Backward.Node>}