src/autograd/backward.bend

bendtorch-core@0.1.0.0 · 337 lines · 46 definitions

1import Base
2
3# A tape is a list of entries, newest first. The entry with k older entries
4# has id k, so ids follow creation order. A lookup walks down from the newest
5# entry and compares each entry's id with the requested one. It never computes
6# an offset by subtraction, so an id the tape does not hold reads as
7# UnknownNode and never as another entry.
8#
9# backward(tape, output) seeds d(output)/d(output) = 1.0 at the selected
10# output, never at the newest entry. Entries newer than the output receive
11# nothing. The sweep then visits the output and every older entry, newest
12# first. An entry that received a gradient sends its vector-Jacobian product
13# to its operands, the first operand before the second; an entry that received
14# nothing sends nothing. A node's first contribution moves in unchanged, and
15# every later one is added with one F32 add, the held total first. PyTorch's
16# engine fills its input buffers in the same order: consumers in decreasing
17# creation order, the operands of one consumer in order.
18#
19# The gradients of one backward are then accumulated into the slot of every
20# entry with the same rule, so a repeated backward adds to the slots like
21# PyTorch's backward(retain_graph=True). The gradients depend only on the
22# nodes, never on the slots. A slot that no backward reached reads None, like
23# PyTorch's .grad. zero_grad empties every slot.
24#
25# The sweep is sequential: the order in which contributions reach a slot is
26# part of the contract.
27
28type Error is Data:
29  UnknownNode{}
30  DanglingOperand{}
31
32type Node is Data:
33  Leaf{}
34  Add{a: Nat, b: Nat}
35  Sub{a: Nat, b: Nat}
36  Mul{a: Nat, b: Nat, left: F32, right: F32}
37  Neg{a: Nat}
38
39type Entry is Data:
40  Entry{node: Node, value: F32, grad: Maybe<&2, F32>}
41
42type Tape is Data:
43  Tape{entries: List<&2, Entry>}
44
45type Contribution is Data:
46  Contribution{target: Nat, value: F32}
47
48type Binary is Data:
49  Plus{}
50  Minus{}
51  Times{}
52
53def empty() -> Tape:
54  Tape{Nil{}}
55
56def size(tape: Tape) -> Nat:
57  Tape{entries} = tape
58  List.length(&2, Entry, entries)
59
60def picked(-A: Data, head: A, rec: Unit -> Maybe<&2, A>, same: Bool) -> Maybe<&2, A>:
61  match same:
62    case True{}:
63      Some{head}
64    case False{}:
65      rec(Unit{})
66
67# The element with id `id` in a newest-first list of count elements: the head
68# has id count - 1, and the walk counts down from it.
69def find(-A: Data, xs: List<&2, A>, +count: Nat, +id: Nat) -> Maybe<&2, A>:
70  match xs count:
71    case Nil{} count:
72      None{}
73    case Con{head, rest} 0n:
74      None{}
75    case Con{head, +rest} 1n+below:
76      picked(A, head, unit => find(A, rest, below, id), Nat.is_eq(below, id))
77
78def lookup(-A: Data, +xs: List<&2, A>, +id: Nat) -> Maybe<&2, A>:
79  find(A, xs, List.length(&2, A, xs), id)
80
81def accumulate(total: Maybe<&2, F32>, contribution: Maybe<&2, F32>) -> Maybe<&2, F32>:
82  match total contribution:
83    case None{} contribution:
84      contribution
85    case Some{held} None{}:
86      Some{held}
87    case Some{held} Some{value}:
88      Some{F32.add(held, value)}
89
90def bumped(slot: Maybe<&2, F32>, +rest: List<&2, Maybe<&2, F32>>, +contribution: F32,
91  rec: Unit -> List<&2, Maybe<&2, F32>>, same: Bool) -> List<&2, Maybe<&2, F32>>:
92  match same:
93    case True{}:
94      accumulate(slot, Some{contribution}) <> rest
95    case False{}:
96      slot <> rec(Unit{})
97
98def bump(slots: List<&2, Maybe<&2, F32>>, +count: Nat, +target: Nat, +contribution: F32) -> List<&2, Maybe<&2, F32>>:
99  match slots count:
100    case Nil{} count:
101      Nil{}
102    case Con{slot, rest} 0n:
103      slot <> rest
104    case Con{slot, +rest} 1n+below:
105      bumped(slot, rest, contribution, unit => bump(rest, below, target, contribution), Nat.is_eq(below, target))
106
107# The vector-Jacobian product of one node: the contribution each operand
108# receives when the node's gradient is g, in operand order.
109def vjp(node: Node, +g: F32) -> List<&2, Contribution>:
110  match node:
111    case Leaf{}:
112      []
113    case Add{a, b}:
114      [Contribution{a, g}, Contribution{b, g}]
115    case Sub{a, b}:
116      [Contribution{a, g}, Contribution{b, F32.neg(g)}]
117    case Mul{a, b, left, right}:
118      [Contribution{a, F32.mul(g, right)}, Contribution{b, F32.mul(g, left)}]
119    case Neg{a}:
120      [Contribution{a, F32.neg(g)}]
121
122def deliver(contributions: List<&2, Contribution>, +count: Nat, slots: List<&2, Maybe<&2, F32>>) -> List<&2, Maybe<&2, F32>>:
123  match contributions:
124    case Nil{}:
125      slots
126    case Con{Contribution{+target, +value}, rest}:
127      deliver(rest, count, bump(slots, count, target, value))
128
129def contribute(node: Node, gradient: Maybe<&2, F32>, +count: Nat, slots: List<&2, Maybe<&2, F32>>) -> List<&2, Maybe<&2, F32>>:
130  match gradient:
131    case None{}:
132      slots
133    case Some{value}:
134      deliver(vjp(node, value), count, slots)
135
136def sweep(nodes: List<&2, Node>, gradients: List<&2, Maybe<&2, F32>>, +count: Nat) -> List<&2, Maybe<&2, F32>>:
137  match nodes gradients count:
138    case Nil{} gradients count:
139      gradients
140    case Con{node, rest} Nil{} count:
141      Nil{}
142    case Con{node, rest} Con{gradient, pending} 0n:
143      Nil{}
144    case Con{node, rest} Con{+gradient, pending} 1n+below:
145      gradient <> sweep(rest, contribute(node, gradient, below, pending), below)
146
147def nones(count: Nat) -> List<&2, Maybe<&2, F32>>:
148  match count:
149    case 0n:
150      Nil{}
151    case 1n+rest:
152      None{} <> nones(rest)
153
154def chosen(node: Node, +rest: List<&2, Node>, +below: Nat, rec: Unit -> List<&2, Maybe<&2, F32>>, selected: Bool) ->
155  List<&2, Maybe<&2, F32>>:
156  match selected:
157    case True{}:
158      sweep(node <> rest, Some{1.0} <> nones(below), 1n+below)
159    case False{}:
160      None{} <> rec(Unit{})
161
162# The gradient of output with respect to each of count nodes, newest first:
163# None for the nodes newer than output and for every node output does not
164# reach.
165def pass(nodes: List<&2, Node>, +count: Nat, +output: Nat) -> List<&2, Maybe<&2, F32>>:
166  match nodes count:
167    case Nil{} count:
168      Nil{}
169    case Con{node, rest} 0n:
170      Nil{}
171    case Con{node, +rest} 1n+below:
172      chosen(node, rest, below, unit => pass(rest, below, output), Nat.is_eq(below, output))
173
174def operands_below(node: Node, +bound: Nat) -> Bool:
175  match node:
176    case Leaf{}:
177      True{}
178    case Add{a, b}:
179      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
180    case Sub{a, b}:
181      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
182    case Mul{a, b, left, right}:
183      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
184    case Neg{a}:
185      Nat.is_lt(a, bound)
186
187def wired(nodes: List<&2, Node>, +count: Nat) -> Bool:
188  match nodes count:
189    case Nil{} count:
190      True{}
191    case Con{node, rest} 0n:
192      False{}
193    case Con{node, rest} 1n+below:
194      Bool.and(operands_below(node, below), wired(rest, below))
195
196def nodes(entries: List<&2, Entry>) -> List<&2, Node>:
197  match entries:
198    case Nil{}:
199      Nil{}
200    case Con{Entry{node, value, slot}, rest}:
201      node <> nodes(rest)
202
203def checked(+graph: List<&2, Node>, +count: Nat, +output: Nat, known: Bool, sound: Bool) ->
204  Result<&2, &2, Error, List<&2, Maybe<&2, F32>>>:
205  match known sound:
206    case False{} sound:
207      Fail{UnknownNode{}}
208    case True{} False{}:
209      Fail{DanglingOperand{}}
210    case True{} True{}:
211      Done{pass(graph, count, output)}
212
213def planned(+graph: List<&2, Node>, +output: Nat) -> Result<&2, &2, Error, List<&2, Maybe<&2, F32>>>:
214  +count = List.length(&2, Node, graph)
215  checked(graph, count, output, Nat.is_lt(output, count), wired(graph, count))
216
217# The gradients of one backward of output, newest first, read from the nodes
218# alone. An output the tape does not hold is refused with UnknownNode; a tape
219# with a node that refers to itself or to a newer or missing node is refused
220# with DanglingOperand.
221def gradients(tape: Tape, +output: Nat) -> Result<&2, &2, Error, List<&2, Maybe<&2, F32>>>:
222  Tape{entries} = tape
223  planned(nodes(entries), output)
224
225def merge(entries: List<&2, Entry>, gradients: List<&2, Maybe<&2, F32>>) -> List<&2, Entry>:
226  match entries gradients:
227    case Nil{} gradients:
228      Nil{}
229    case Con{entry, rest} Nil{}:
230      entry <> rest
231    case Con{Entry{node, value, slot}, rest} Con{gradient, more}:
232      Entry{node, value, accumulate(slot, gradient)} <> merge(rest, more)
233
234def merged(entries: List<&2, Entry>, result: Result<&2, &2, Error, List<&2, Maybe<&2, F32>>>) -> Result<&2, &2, Error, Tape>:
235  match result:
236    case Fail{error}:
237      Fail{error}
238    case Done{found}:
239      Done{Tape{merge(entries, found)}}
240
241# Runs one backward of output and accumulates its gradients into the slots.
242def backward(tape: Tape, +output: Nat) -> Result<&2, &2, Error, Tape>:
243  Tape{+entries} = tape
244  merged(entries, gradients(Tape{entries}, output))
245
246def cleared(entries: List<&2, Entry>) -> List<&2, Entry>:
247  match entries:
248    case Nil{}:
249      Nil{}
250    case Con{Entry{node, value, slot}, rest}:
251      Entry{node, value, None{}} <> cleared(rest)
252
253# Empties every gradient slot, like PyTorch's zero_grad(set_to_none=True).
254def zero_grad(tape: Tape) -> Tape:
255  Tape{entries} = tape
256  Tape{cleared(entries)}
257
258def entry_value(found: Maybe<&2, Entry>) -> Result<&2, &2, Error, F32>:
259  match found:
260    case None{}:
261      Fail{UnknownNode{}}
262    case Some{Entry{node, value, slot}}:
263      Done{value}
264
265def entry_grad(found: Maybe<&2, Entry>) -> Result<&2, &2, Error, Maybe<&2, F32>>:
266  match found:
267    case None{}:
268      Fail{UnknownNode{}}
269    case Some{Entry{node, value, slot}}:
270      Done{slot}
271
272def value(tape: Tape, +id: Nat) -> Result<&2, &2, Error, F32>:
273  Tape{entries} = tape
274  entry_value(lookup(Entry, entries, id))
275
276# The gradient slot of node id: None until a backward reaches it.
277def grad(tape: Tape, +id: Nat) -> Result<&2, &2, Error, Maybe<&2, F32>>:
278  Tape{entries} = tape
279  entry_grad(lookup(Entry, entries, id))
280
281# Records a leaf; its id is the number of entries before it.
282def leaf(tape: Tape, value: F32) -> Tape & Nat:
283  Tape{+entries} = tape
284  (Tape{Entry{Leaf{}, value, None{}} <> entries}, List.length(&2, Entry, entries))
285
286def forward(kind: Binary, +a: Nat, +b: Nat, +left: F32, +right: F32) -> Entry:
287  match kind:
288    case Plus{}:
289      Entry{Add{a, b}, F32.add(left, right), None{}}
290    case Minus{}:
291      Entry{Sub{a, b}, F32.sub(left, right), None{}}
292    case Times{}:
293      Entry{Mul{a, b, left, right}, F32.mul(left, right), None{}}
294
295def operands(+entries: List<&2, Entry>, kind: Binary, +a: Nat, +b: Nat,
296  left: Result<&2, &2, Error, F32>, right: Result<&2, &2, Error, F32>) -> Tape & Result<&2, &2, Error, Nat>:
297  match left right:
298    case Fail{error} right:
299      (Tape{entries}, Fail{UnknownNode{}})
300    case Done{x} Fail{error}:
301      (Tape{entries}, Fail{UnknownNode{}})
302    case Done{x} Done{y}:
303      (Tape{forward(kind, a, b, x, y) <> entries}, Done{List.length(&2, Entry, entries)})
304
305# Records `a op b`. An operand id the tape does not hold is refused with
306# UnknownNode, and the tape comes back unchanged.
307def binary(tape: Tape, kind: Binary, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
308  Tape{+entries} = tape
309  operands(entries, kind, a, b, value(Tape{entries}, a), value(Tape{entries}, b))
310
311def add(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
312  binary(tape, Plus{}, a, b)
313
314def sub(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
315  binary(tape, Minus{}, a, b)
316
317def mul(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
318  binary(tape, Times{}, a, b)
319
320def negated(+entries: List<&2, Entry>, +a: Nat, operand: Result<&2, &2, Error, F32>) -> Tape & Result<&2, &2, Error, Nat>:
321  match operand:
322    case Fail{error}:
323      (Tape{entries}, Fail{UnknownNode{}})
324    case Done{x}:
325      (Tape{Entry{Neg{a}, F32.neg(x), None{}} <> entries}, Done{List.length(&2, Entry, entries)})
326
327# Records -a, with the same refusal as binary.
328def neg(tape: Tape, +a: Nat) -> Tape & Result<&2, &2, Error, Nat>:
329  Tape{+entries} = tape
330  negated(entries, a, value(Tape{entries}, a))
331
332def error_code(error: Error) -> U32:
333  match error:
334    case UnknownNode{}:
335      103
336    case DanglingOperand{}:
337      104