1	import 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	
28	type Error is Data:
29	  UnknownNode{}
30	  DanglingOperand{}
31	
32	type 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	
39	type Entry is Data:
40	  Entry{node: Node, value: F32, grad: Maybe<&2, F32>}
41	
42	type Tape is Data:
43	  Tape{entries: List<&2, Entry>}
44	
45	type Contribution is Data:
46	  Contribution{target: Nat, value: F32}
47	
48	type Binary is Data:
49	  Plus{}
50	  Minus{}
51	  Times{}
52	
53	def empty() -> Tape:
54	  Tape{Nil{}}
55	
56	def size(tape: Tape) -> Nat:
57	  Tape{entries} = tape
58	  List.length(&2, Entry, entries)
59	
60	def 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.
69	def 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	
78	def lookup(-A: Data, +xs: List<&2, A>, +id: Nat) -> Maybe<&2, A>:
79	  find(A, xs, List.length(&2, A, xs), id)
80	
81	def 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	
90	def 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	
98	def 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.
109	def 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	
122	def 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	
129	def 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	
136	def 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	
147	def 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	
154	def 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.
165	def 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	
174	def 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	
187	def 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	
196	def 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	
203	def 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	
213	def 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.
221	def gradients(tape: Tape, +output: Nat) -> Result<&2, &2, Error, List<&2, Maybe<&2, F32>>>:
222	  Tape{entries} = tape
223	  planned(nodes(entries), output)
224	
225	def 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	
234	def 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.
242	def backward(tape: Tape, +output: Nat) -> Result<&2, &2, Error, Tape>:
243	  Tape{+entries} = tape
244	  merged(entries, gradients(Tape{entries}, output))
245	
246	def 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).
254	def zero_grad(tape: Tape) -> Tape:
255	  Tape{entries} = tape
256	  Tape{cleared(entries)}
257	
258	def 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	
265	def 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	
272	def 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.
277	def 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.
282	def leaf(tape: Tape, value: F32) -> Tape & Nat:
283	  Tape{+entries} = tape
284	  (Tape{Entry{Leaf{}, value, None{}} <> entries}, List.length(&2, Entry, entries))
285	
286	def 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	
295	def 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.
307	def 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	
311	def add(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
312	  binary(tape, Plus{}, a, b)
313	
314	def sub(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
315	  binary(tape, Minus{}, a, b)
316	
317	def mul(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
318	  binary(tape, Times{}, a, b)
319	
320	def 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.
328	def neg(tape: Tape, +a: Nat) -> Tape & Result<&2, &2, Error, Nat>:
329	  Tape{+entries} = tape
330	  negated(entries, a, value(Tape{entries}, a))
331	
332	def error_code(error: Error) -> U32:
333	  match error:
334	    case UnknownNode{}:
335	      103
336	    case DanglingOperand{}:
337	      104