src/autograd/tensor_backward.bend

bendtorch-core@0.1.0.0 · 1,698 lines · 215 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/elementwise.bend as Elementwise
5import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
6import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
7import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as Batched
8import bendtorch-core@0.1.0.0/src/tensor/reduction.bend as Reduction
9import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
10import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
11import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
12import bendtorch-core@0.1.0.0/src/tensor/gather.bend as Gather
13import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
14import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
15import bendtorch-core@0.1.0.0/src/autograd/sum_to_size.bend as SumToSize
16import bendtorch-core@0.1.0.0/src/autograd/sum_dims.bend as SumDims
17import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
18import bendtorch-core@0.1.0.0/src/autograd/loss.bend as Loss
19
20# Reverse mode over tensors, on the scalar engine's tape discipline
21# (backward.bend): entries newest first, the entry with k older entries has
22# id k, a lookup compares ids and never subtracts, the seed goes to the
23# selected output, and the sweep visits the output and every older entry,
24# newest first. An entry that received a gradient sends its vector-Jacobian
25# products to its operands, the first operand before the second; an entry that
26# received nothing sends nothing.
27#
28# Each contribution is reduced to its operand's shape with sum_to_size, as
29# PyTorch's engine does for a broadcast operand, and accumulated into the
30# operand's working slot with Gradient.accumulate: the first contribution
31# moves in unchanged, every later one is added position by position, the held
32# total first. After the sweep, each working gradient is accumulated into its
33# entry's slot with the same rule, so a repeated backward adds, like PyTorch's
34# backward(retain_graph=True).
35#
36# The seed is a tensor of ones with the output's shape, as PyTorch's
37# backward(torch.ones_like(output)). The products are
38#   Add: g and g; Sub: g and -g; Mul: g * right and g * left, with the operand
39#   values saved at record time; Sum: g broadcast to the operand's shape;
40#   Matmul: g @ right^T and left^T @ g; MatmulNT (left @ right^T, the
41#   product of Linear): g @ right and g^T @ left, read in place; Relu: g where
42#   the saved input is above zero or NaN and +0.0 where it is at most zero,
43#   PyTorch's threshold_backward; MseLoss: (2 / N * d) * g and (2 / N * e) * g
44#   with the saved differences d = left - right and e = right - left;
45#   CrossEntropy: two contributions to the logits, one_hot * -(g / N) and
46#   -(exp(lp) * (-(g / N) + 0.0)) with the saved log_softmax lp, whose sum is
47#   PyTorch's one_hot gradient minus exp(lp) times its row sum;
48#   BatchedMatmul (torch.matmul of operands that are not both rank 2): g @
49#   right^T and left^T @ g over the output's batch, each operand read in
50#   place at its own batch, 1-D operands as a row or a column;
51#   Slice (a slice, narrow or select): g read with the kept shape, the
52#   selected axis put back as 1 for select, then scattered into +0.0 words of
53#   the operand's shape, PyTorch's slice_backward.
54# Each is reduced to its operand's shape before it is delivered.
55#
56# A constant is a leaf that takes no gradient, PyTorch's tensor with
57# requires_grad=False. A recording whose operands are all constants records a
58# constant too, as PyTorch's result does not require grad either. Any other
59# recording that reads a constant wraps its node in Frozen{id, node}, one wrap
60# for each operand id that is a constant, and the VJP of a frozen node sends
61# nothing to id: it drops every contribution to id from node's VJP, so a
62# constant's slot stays empty, as PyTorch's .grad stays None. Where the
63# dropped contribution is the product g @ W of `matmul_nt`, the input of
64# Linear, the VJP does not compute it.
65#
66# Recording and reads hand back the tape. A failed call leaves every slot as
67# it was.
68
69type Error is Data:
70  UnknownNode{}
71  DanglingOperand{}
72  TensorError{error: Core.Error}
73  MatrixError{error: Matrix.Error}
74  TransformError{error: Transforms.Error}
75  BroadcastError{error: Broadcast.Error}
76  ReduceError{error: Reduce.Error}
77  GradientError{error: Gradient.Error}
78  LossError{error: Loss.Error}
79  SliceError{error: Slice.Error}
80  GatherError{error: Gather.Error}
81
82type Node is Data:
83  Leaf{}
84  Add{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
85  Sub{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
86  Mul{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
87  Sum{a: Nat, input: Shape.Shape}
88  SumDims{a: Nat, input: Shape.Shape, kept: Shape.Shape}
89  Matmul{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
90  MatmulNT{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
91  Relu{a: Nat, input: Shape.Shape}
92  MseLoss{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
93  CrossEntropy{a: Nat, input: Shape.Shape}
94  BatchedMatmul{a: Nat, b: Nat, left: Shape.Shape, right: Shape.Shape}
95  Slice{a: Nat, input: Shape.Shape, kept: Shape.Shape, plan: Slice.Plan}
96  Gathered{a: Nat, input: Shape.Shape, plan: Gather.Plan, entries: List<&2, U32>}
97  Constant{}
98  Frozen{id: Nat, node: Node}
99
100# The operand values a product saves at record time.
101type Saved is Type:
102  Unsaved{}
103  Operands{left: Core.Dynamic, right: Core.Dynamic}
104
105type Entry is Type:
106  Entry{node: Node, shape: Shape.Shape, value: Core.Dynamic, saved: Saved, grad: Gradient.Slot}
107
108type Tape is Type:
109  Tape{entries: List<&1, Entry>}
110
111type Contribution is Type:
112  Contribution{target: Nat, shape: Shape.Shape, gradient: Core.Dynamic}
113
114# What one entry's VJP hands back: its saved values and its gradient, unchanged,
115# and the contributions for its operands in operand order.
116type Pass is Type:
117  Pass{saved: Saved, gradient: Core.Dynamic, contributions: Result<&2, &1, Error, List<&1, Contribution>>}
118
119type Binary is Data:
120  Plus{}
121  Minus{}
122  Times{}
123  Product{}
124  ProductNT{}
125  SquaredError{}
126  BatchedProduct{}
127
128def empty() -> Tape:
129  Tape{Nil{}}
130
131# Whether the head of a newest-first list of count elements has id `id`.
132def at_head(+count: Nat, +id: Nat) -> Bool:
133  match count:
134    case 0n:
135      False{}
136    case 1n+below:
137      Nat.is_eq(below, id)
138
139def entry_listed(entry: Entry, pair: List<&1, Entry> & List<&2, Node>) -> List<&1, Entry> & List<&2, Node>:
140  Entry{+node, shape, value, saved, slot} = entry
141  (rest, nodes) = pair
142  (Entry{node, shape, value, saved, slot} <> rest, node <> nodes)
143
144# The nodes of a list of entries, newest first, next to the entries.
145def listed(entries: List<&1, Entry>) -> List<&1, Entry> & List<&2, Node>:
146  match entries:
147    case Nil{}:
148      (Nil{}, Nil{})
149    case Con{entry, rest}:
150      entry_listed(entry, listed(rest))
151
152def prepend(-A: Type, head: A, result: Result<&2, &1, Error, List<&1, A>>) -> Result<&2, &1, Error, List<&1, A>>:
153  match result:
154    case Fail{error}:
155      Fail{error}
156    case Done{rest}:
157      Done{head <> rest}
158
159def unsets(count: Nat) -> List<&1, Gradient.Slot>:
160  match count:
161    case 0n:
162      Nil{}
163    case 1n+rest:
164      Gradient.Unset{} <> unsets(rest)
165
166# Vector-Jacobian products.
167
168def both(first: Result<&2, &1, Error, Contribution>, second: Result<&2, &1, Error, Contribution>) ->
169  Result<&2, &1, Error, List<&1, Contribution>>:
170  match first second:
171    case Fail{error} second:
172      Fail{error}
173    case Done{left} Fail{error}:
174      Fail{error}
175    case Done{left} Done{right}:
176      Done{[left, right]}
177
178def delivered_as(+target: Nat, +shape: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) ->
179  Core.Dynamic & Result<&2, &1, Error, Contribution>:
180  (owner, result) = pair
181  match result:
182    case Fail{error}:
183      (owner, Fail{ReduceError{error}})
184    case Done{reduced}:
185      (owner, Done{Contribution{target, shape, reduced}})
186
187# A contribution for operand target: g reduced to the operand's shape. g comes
188# back unchanged.
189def reduced(+target: Nat, +shape: Shape.Shape, g: Core.Dynamic) -> Core.Dynamic & Result<&2, &1, Error, Contribution>:
190  delivered_as(target, shape, SumToSize.sum_to_size(g, shape))
191
192def contribution_of(pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Result<&2, &1, Error, Contribution>:
193  (owner, result) = pair
194  result
195
196def product_reduced(+target: Nat, +shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Dynamic>) ->
197  Result<&2, &1, Error, Contribution>:
198  match result:
199    case Fail{error}:
200      Fail{TensorError{error}}
201    case Done{product}:
202      contribution_of(reduced(target, shape, product))
203
204def added_second(first: Result<&2, &1, Error, Contribution>, pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
205  (g, second) = pair
206  Pass{Unsaved{}, g, both(first, second)}
207
208def added_first(+b: Nat, +right: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
209  (g, first) = pair
210  added_second(first, reduced(b, right, g))
211
212# d(a+b)/da and d(a+b)/db: g, each reduced to its operand's shape.
213def vjp_add(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, g: Core.Dynamic) -> Pass:
214  added_first(b, right, reduced(a, left, g))
215
216def negated_pair(pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
217  Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
218  (operands, result) = pair
219  (g, zero) = operands
220  (g, result)
221
222# -g, word by word: F32.neg flips the sign bit, as PyTorch's neg does.
223def negated(g: Core.Dynamic, zero: Result<&2, &1, Core.Error, Core.Tensor<Reduction.scalar_shape()>>) ->
224  Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
225  match zero:
226    case Fail{error}:
227      (g, Fail{error})
228    case Done{tensor}:
229      negated_pair(Elementwise.binary(Elementwise.NegateLeft{}, g, Core.pack(Reduction.scalar_shape(), tensor)))
230
231def negate(g: Core.Dynamic) -> Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
232  negated(g, Core.new(Reduction.scalar_shape(), 0.0))
233
234def subtracted_second(first: Result<&2, &1, Error, Contribution>, +b: Nat, +right: Shape.Shape,
235  pair: Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
236  (g, minus) = pair
237  Pass{Unsaved{}, g, both(first, product_reduced(b, right, minus))}
238
239def subtracted_first(+b: Nat, +right: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
240  (g, first) = pair
241  subtracted_second(first, b, right, negate(g))
242
243# d(a-b)/da: g; d(a-b)/db: -g, negated before it is reduced, as in PyTorch.
244def vjp_sub(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, g: Core.Dynamic) -> Pass:
245  subtracted_first(b, right, reduced(a, left, g))
246
247def multiplied_second(+b: Nat, +right: Shape.Shape, first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic,
248  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
249  (operands, second) = pair
250  (g, x) = operands
251  Pass{Operands{x, y}, g, both(first, product_reduced(b, right, second))}
252
253def multiplied_first(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, x: Core.Dynamic,
254  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
255  (operands, first) = pair
256  (g, y) = operands
257  multiplied_second(b, right, product_reduced(a, left, first), y, Elementwise.mul(g, x))
258
259# d(a*b)/da: g * right; d(a*b)/db: g * left, with the operand values saved
260# at record time, each product reduced to its operand's shape.
261def vjp_mul(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
262  match saved:
263    case Unsaved{}:
264      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
265    case Operands{x, y}:
266      multiplied_first(a, b, left, right, x, Elementwise.mul(g, y))
267
268def expanded(+a: Nat, +input: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Pass:
269  (g, result) = pair
270  match result:
271    case Fail{error}:
272      Pass{Unsaved{}, g, Fail{BroadcastError{error}}}
273    case Done{expanded}:
274      Pass{Unsaved{}, g, Done{[Contribution{a, input, expanded}]}}
275
276# d(sum a)/da: g broadcast to a's shape, as PyTorch's g.expand(a.shape).
277def vjp_sum(+a: Nat, +input: Shape.Shape, g: Core.Dynamic) -> Pass:
278  expanded(a, input, Broadcast.broadcast_to(g, input))
279
280def relabel_decided(+held: Shape.Shape, +shape: Shape.Shape, tensor: Core.Tensor<held>,
281  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(held), Shape.elements(shape))>) -> Core.Dynamic:
282  match decision:
283    case Shape.BFalse{equal}:
284      Core.Packed{held, tensor}
285    case Shape.BTrue{equal}:
286      Transforms.reshape_equal(held, shape, tensor, Transforms.elements_equal(held, shape, equal))
287
288# g with the dimensions of `shape`, which has as many elements: a reshape,
289# which moves no word. A g with another element count keeps its own.
290def relabeled(+shape: Shape.Shape, gradient: Core.Dynamic) -> Core.Dynamic:
291  Core.Packed{+held, tensor} = gradient
292  relabel_decided(held, shape, tensor, Shape.bool_case(Nat.is_eq(Shape.elements(held), Shape.elements(shape))))
293
294# The pass with g back in its own shape.
295def restored(+shape: Shape.Shape, pass: Pass) -> Pass:
296  Pass{saved, gradient, contributions} = pass
297  Pass{saved, relabeled(shape, gradient), contributions}
298
299def unsqueezed(+a: Nat, +input: Shape.Shape, +kept: Shape.Shape, +shape: Shape.Shape, tensor: Core.Tensor<shape>,
300  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(shape), Shape.elements(kept))>) -> Pass:
301  match decision:
302    case Shape.BFalse{equal}:
303      Pass{Unsaved{}, Core.Packed{shape, tensor}, Fail{TransformError{Transforms.LayoutError{Layout.ElementCountMismatch{}}}}}
304    case Shape.BTrue{equal}:
305      restored(shape, vjp_sum(a, input, Transforms.reshape_equal(shape, kept, tensor, Transforms.elements_equal(shape, kept, equal))))
306
307# d(sum of a over some axes)/da: g read with the kept shape, the reduced axes
308# put back as 1 (PyTorch's unsqueeze when keepdim is False), then broadcast
309# to a's shape, as the Sum VJP does.
310def vjp_sum_dims(+a: Nat, +input: Shape.Shape, +kept: Shape.Shape, g: Core.Dynamic) -> Pass:
311  Core.Packed{+shape, tensor} = g
312  unsqueezed(a, input, kept, shape, tensor, Shape.bool_case(Nat.is_eq(Shape.elements(shape), Shape.elements(kept))))
313
314def multiplied(+target: Nat, +shape: Shape.Shape,
315  result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Result<&2, &1, Error, Contribution>:
316  match result:
317    case Fail{error}:
318      Fail{MatrixError{error}}
319    case Done{triple}:
320      (left, right, product) = triple
321      Done{Contribution{target, shape, product}}
322
323# g @ right^T for the left operand. g and right are copied first, so a failure
324# loses neither.
325def left_product(+target: Nat, +shape: Shape.Shape, copy: Core.Dynamic,
326  pair: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Core.Dynamic & Result<&2, &1, Error, Contribution>:
327  (right, transposed) = pair
328  match transposed:
329    case Fail{error}:
330      (right, Fail{TransformError{error}})
331    case Done{flipped}:
332      (right, multiplied(target, shape, Matrix.matmul(copy, flipped)))
333
334# left^T @ g for the right operand.
335def right_product(+target: Nat, +shape: Shape.Shape, copy: Core.Dynamic,
336  pair: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Core.Dynamic & Result<&2, &1, Error, Contribution>:
337  (left, transposed) = pair
338  match transposed:
339    case Fail{error}:
340      (left, Fail{TransformError{error}})
341    case Done{flipped}:
342      (left, multiplied(target, shape, Matrix.matmul(flipped, copy)))
343
344def matmul_fourth(first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic, g: Core.Dynamic,
345  pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
346  (x, second) = pair
347  Pass{Operands{x, y}, g, both(first, second)}
348
349def matmul_third(+b: Nat, +right: Shape.Shape, first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic, x: Core.Dynamic,
350  pair: Core.Dynamic & Core.Dynamic) -> Pass:
351  (g, copy) = pair
352  matmul_fourth(first, y, g, right_product(b, right, copy, Transforms.transpose_dynamic(x, 0, 1)))
353
354def matmul_second(+b: Nat, +right: Shape.Shape, x: Core.Dynamic, g: Core.Dynamic,
355  pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
356  (y, first) = pair
357  matmul_third(b, right, first, y, x, Core.clone_dynamic(g))
358
359def matmul_first(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, x: Core.Dynamic, y: Core.Dynamic,
360  pair: Core.Dynamic & Core.Dynamic) -> Pass:
361  (g, copy) = pair
362  matmul_second(b, right, x, g, left_product(a, left, copy, Transforms.transpose_dynamic(y, 0, 1)))
363
364# d(a@b)/da: g @ b^T; d(a@b)/db: a^T @ g, with the operand values saved at
365# record time.
366def vjp_matmul(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
367  match saved:
368    case Unsaved{}:
369      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
370    case Operands{x, y}:
371      matmul_first(a, b, left, right, x, y, Core.clone_dynamic(g))
372
373# g @ W for the left operand of x @ W^T: g and W are copied first, so a
374# failure loses neither.
375def nt_left(+target: Nat, +shape: Shape.Shape, copy: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) ->
376  Core.Dynamic & Result<&2, &1, Error, Contribution>:
377  (right, flipped) = pair
378  (right, multiplied(target, shape, Matrix.matmul(copy, flipped)))
379
380# g^T @ x, g read in place. It has its own name so that the proofs of where
381# the contribution goes name it without restating the product; the value
382# laws read what it computes.
383def weight_product(copy: Core.Dynamic, flipped: Core.Dynamic) -> Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
384  Matrix.matmul_tn(copy, flipped)
385
386# g^T @ x for the right operand.
387def nt_right(+target: Nat, +shape: Shape.Shape, copy: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) ->
388  Core.Dynamic & Result<&2, &1, Error, Contribution>:
389  (left, flipped) = pair
390  (left, multiplied(target, shape, weight_product(copy, flipped)))
391
392def nt_fourth(first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic, g: Core.Dynamic,
393  pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
394  (x, second) = pair
395  Pass{Operands{x, y}, g, both(first, second)}
396
397def nt_third(+b: Nat, +right: Shape.Shape, first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic, x: Core.Dynamic,
398  pair: Core.Dynamic & Core.Dynamic) -> Pass:
399  (g, copy) = pair
400  nt_fourth(first, y, g, nt_right(b, right, copy, Core.clone_dynamic(x)))
401
402def nt_second(+b: Nat, +right: Shape.Shape, x: Core.Dynamic, g: Core.Dynamic,
403  pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
404  (y, first) = pair
405  nt_third(b, right, first, y, x, Core.clone_dynamic(g))
406
407def nt_first(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, x: Core.Dynamic, y: Core.Dynamic,
408  pair: Core.Dynamic & Core.Dynamic) -> Pass:
409  (g, copy) = pair
410  nt_second(b, right, x, g, nt_left(a, left, copy, Core.clone_dynamic(y)))
411
412def alone(result: Result<&2, &1, Error, Contribution>) -> Result<&2, &1, Error, List<&1, Contribution>>:
413  match result:
414    case Fail{error}:
415      Fail{error}
416    case Done{contribution}:
417      Done{[contribution]}
418
419def nt_alone_second(y: Core.Dynamic, g: Core.Dynamic, pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
420  (x, second) = pair
421  Pass{Operands{x, y}, g, alone(second)}
422
423def nt_alone_first(+b: Nat, +right: Shape.Shape, x: Core.Dynamic, y: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) -> Pass:
424  (g, copy) = pair
425  nt_alone_second(y, g, nt_right(b, right, copy, Core.clone_dynamic(x)))
426
427# d(x @ W^T)/dW alone, g^T @ x, for a product whose input x takes no
428# gradient: the second contribution of vjp_matmul_nt, computed the same way,
429# without g @ W.
430def vjp_matmul_nt_right(+b: Nat, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
431  match saved:
432    case Unsaved{}:
433      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
434    case Operands{x, y}:
435      nt_alone_first(b, right, x, y, Core.clone_dynamic(g))
436
437# d(x @ W^T)/dx: g @ W; d(x @ W^T)/dW: g^T @ x, with the operand values
438# saved at record time. Neither product builds a transpose: g @ W is
439# `matmul`, and g^T @ x is `matmul_tn`, which reads g in place. Each product
440# already has its operand's shape.
441def vjp_matmul_nt(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
442  match saved:
443    case Unsaved{}:
444      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
445    case Operands{x, y}:
446      nt_first(a, b, left, right, x, y, Core.clone_dynamic(g))
447
448def relu_passed(+a: Nat, +input: Shape.Shape, zero: Core.Dynamic,
449  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
450  (operands, result) = pair
451  (x, g) = operands
452  Pass{Operands{x, zero}, g, alone(product_reduced(a, input, result))}
453
454# d relu(a)/da: the pointwise ReluBackward of the saved input x and g, g's
455# own word where x is above zero or NaN and +0.0 where x is at most zero,
456# reduced to the input's shape.
457def vjp_relu(+a: Nat, +input: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
458  match saved:
459    case Unsaved{}:
460      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
461    case Operands{x, zero}:
462      relu_passed(a, input, zero, Elementwise.binary(Elementwise.ReluBackward{}, x, g))
463
464def mse_second(+b: Nat, +right: Shape.Shape, first: Result<&2, &1, Error, Contribution>, d: Core.Dynamic,
465  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
466  (operands, second) = pair
467  (g, e) = operands
468  Pass{Operands{e, d}, g, both(first, product_reduced(b, right, second))}
469
470def mse_first(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +scale: F32, e: Core.Dynamic,
471  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
472  (operands, first) = pair
473  (g, d) = operands
474  mse_second(b, right, product_reduced(a, left, first), d, Elementwise.binary(Elementwise.MseGrad{scale}, g, e))
475
476def mse_shaped(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, e: Core.Dynamic, g: Core.Dynamic,
477  pair: Core.Dynamic & Shape.Shape) -> Pass:
478  (d, +shape) = pair
479  mse_first(a, b, left, right, Loss.mse_scale(shape), e, Elementwise.binary(Elementwise.MseGrad{Loss.mse_scale(shape)}, g, d))
480
481# d mse_loss(a, b)/da: (2 / N * d) * g, and d/db: (2 / N * e) * g, PyTorch's
482# mse_loss_backward, with the differences e = right - left and d = left -
483# right saved at record time over the broadcast shape of N cells; each is
484# reduced to its operand's shape. The steps are the Mul VJP's, with
485# MseGrad for Mul.
486def vjp_mse(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
487  match saved:
488    case Unsaved{}:
489      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
490    case Operands{e, d}:
491      mse_shaped(a, b, left, right, e, g, Core.shape(d))
492
493def entropy_second(+a: Nat, +input: Shape.Shape, first: Result<&2, &1, Error, Contribution>, onehot: Core.Dynamic,
494  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
495  (operands, second) = pair
496  (g, lp) = operands
497  Pass{Operands{lp, onehot}, g, both(first, product_reduced(a, input, second))}
498
499def entropy_first(+a: Nat, +input: Shape.Shape, +count: F32, +single: Bool, lp: Core.Dynamic,
500  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Pass:
501  (operands, first) = pair
502  (g, onehot) = operands
503  entropy_second(a, input, product_reduced(a, input, first), onehot, Elementwise.binary(Elementwise.SoftmaxGrad{count, single}, g, lp))
504
505# d cross_entropy(a)/da, as two contributions to a that its slot adds in
506# order. The first is PyTorch's nll_loss backward, OneHotGrad of g and the
507# saved one-hot rows: -(g / N) at each row's target and +0.0 elsewhere, for
508# the N rows. The second is the rest of PyTorch's log_softmax backward,
509# SoftmaxGrad of g and the saved lp: -(exp(lp) * r) with r the row sum of the
510# first, -(g / N) + 0.0, or -(g / N) itself when the logits have one column
511# (PyTorch's sum of one word keeps -0.0). Added into an empty slot, first + second is
512# PyTorch's first - exp(lp) * r word for word. The steps are the Mul VJP's,
513# with OneHotGrad and SoftmaxGrad for Mul and both contributions to a.
514def vjp_cross_entropy(+a: Nat, +input: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
515  match saved:
516    case Unsaved{}:
517      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
518    case Operands{lp, onehot}:
519      entropy_first(a, input, Loss.count_f32(Loss.rows(input)), Loss.one_class(input), lp,
520        Elementwise.binary(Elementwise.OneHotGrad{Loss.count_f32(Loss.rows(input))}, g, onehot))
521
522def kept_one(+target: Nat, +shape: Shape.Shape, gradient: Core.Dynamic, mine: Bool, rest: List<&1, Contribution>) ->
523  List<&1, Contribution>:
524  match mine:
525    case True{}:
526      rest
527    case False{}:
528      Contribution{target, shape, gradient} <> rest
529
530# The contributions to every node but id, in order.
531def kept_others(+id: Nat, items: List<&1, Contribution>) -> List<&1, Contribution>:
532  match items:
533    case Nil{}:
534      Nil{}
535    case Con{Contribution{+target, +shape, gradient}, rest}:
536      kept_one(target, shape, gradient, Nat.is_eq(id, target), kept_others(id, rest))
537
538def without(+id: Nat, contributions: Result<&2, &1, Error, List<&1, Contribution>>) -> Result<&2, &1, Error, List<&1, Contribution>>:
539  match contributions:
540    case Fail{error}:
541      Fail{error}
542    case Done{items}:
543      Done{kept_others(id, items)}
544
545# A VJP's pass without its contributions to id.
546def unsent(+id: Nat, pass: Pass) -> Pass:
547  Pass{saved, g, contributions} = pass
548  Pass{saved, g, without(id, contributions)}
549
550# Whether the VJP of Frozen{id, node} skips a product instead of dropping it:
551# node is x @ W^T and id is x and not W.
552def skips(+id: Nat, node: Node) -> Bool:
553  match node:
554    case MatmulNT{a, b, left, right}:
555      Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))
556    case other:
557      False{}
558
559def nt_frozen(+id: Nat, +a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic, skip: Bool) -> Pass:
560  match skip:
561    case True{}:
562      vjp_matmul_nt_right(b, right, saved, g)
563    case False{}:
564      unsent(id, vjp_matmul_nt(a, b, left, right, saved, g))
565
566# A contribution for operand target from a product over the output's batch:
567# the product reduced to the operand's shape, as a broadcast operand of an
568# elementwise operation is.
569def batched_reduced(+target: Nat, +shape: Shape.Shape,
570  result: Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Result<&2, &1, Error, Contribution>:
571  match result:
572    case Fail{error}:
573      Fail{TensorError{error}}
574    case Done{triple}:
575      (first, second, product) = triple
576      contribution_of(reduced(target, shape, product))
577
578def batched_fourth(first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic, g: Core.Dynamic,
579  pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
580  (x, second) = pair
581  Pass{Operands{x, y}, g, both(first, second)}
582
583# left^T @ g for the right operand, from copies of g and of the saved left
584# operand.
585def batched_right(+b: Nat, +right: Shape.Shape, +plan: Batched.Plan, copy: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) ->
586  Core.Dynamic & Result<&2, &1, Error, Contribution>:
587  (x, left_copy) = pair
588  (x, batched_reduced(b, right, Batched.right_gradient(plan, left_copy, copy)))
589
590def batched_third(+b: Nat, +right: Shape.Shape, +plan: Batched.Plan, first: Result<&2, &1, Error, Contribution>, y: Core.Dynamic,
591  x: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) -> Pass:
592  (g, copy) = pair
593  batched_fourth(first, y, g, batched_right(b, right, plan, copy, Core.clone_dynamic(x)))
594
595def batched_second(+b: Nat, +right: Shape.Shape, +plan: Batched.Plan, x: Core.Dynamic, g: Core.Dynamic,
596  pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Pass:
597  (y, first) = pair
598  batched_third(b, right, plan, first, y, x, Core.clone_dynamic(g))
599
600# g @ right^T for the left operand, from copies of g and of the saved right
601# operand, so a failure loses neither.
602def batched_left(+a: Nat, +left: Shape.Shape, +plan: Batched.Plan, copy: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) ->
603  Core.Dynamic & Result<&2, &1, Error, Contribution>:
604  (y, right_copy) = pair
605  (y, batched_reduced(a, left, Batched.left_gradient(plan, copy, right_copy)))
606
607def batched_first(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, +plan: Batched.Plan, x: Core.Dynamic, y: Core.Dynamic,
608  pair: Core.Dynamic & Core.Dynamic) -> Pass:
609  (g, copy) = pair
610  batched_second(b, right, plan, x, g, batched_left(a, left, plan, copy, Core.clone_dynamic(y)))
611
612def batched_planned(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, x: Core.Dynamic, y: Core.Dynamic, g: Core.Dynamic,
613  result: Result<&2, &2, Matrix.Error, Batched.Plan>) -> Pass:
614  match result:
615    case Fail{error}:
616      Pass{Operands{x, y}, g, Fail{MatrixError{error}}}
617    case Done{+plan}:
618      batched_first(a, b, left, right, plan, x, y, Core.clone_dynamic(g))
619
620# d(a@b)/da and d(a@b)/db of torch.matmul: g @ b^T and a^T @ g over the
621# output's batch (Batched.left_gradient and Batched.right_gradient), each
622# then summed to its operand's shape with sum_to_size, which sums the batch
623# dimensions the operand was broadcast along and keeps a product that
624# already has the operand's shape. A 1-D operand's product already drops its
625# row or column.
626def vjp_batched(+a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
627  match saved:
628    case Unsaved{}:
629      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
630    case Operands{x, y}:
631      batched_planned(a, b, left, right, x, y, g, Batched.plan(left, right))
632
633def scatter_sent(+a: Nat, +input: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>) -> Pass:
634  (g, result) = pair
635  match result:
636    case Fail{error}:
637      Pass{Unsaved{}, g, Fail{SliceError{error}}}
638    case Done{scattered}:
639      Pass{Unsaved{}, g, Done{[Contribution{a, input, scattered}]}}
640
641def slice_unsqueezed(+a: Nat, +input: Shape.Shape, +kept: Shape.Shape, +plan: Slice.Plan, +shape: Shape.Shape, tensor: Core.Tensor<shape>,
642  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(shape), Shape.elements(kept))>) -> Pass:
643  match decision:
644    case Shape.BFalse{equal}:
645      Pass{Unsaved{}, Core.Packed{shape, tensor}, Fail{TransformError{Transforms.LayoutError{Layout.ElementCountMismatch{}}}}}
646    case Shape.BTrue{equal}:
647      restored(shape, scatter_sent(a, input,
648        Slice.scatter(Transforms.reshape_equal(shape, kept, tensor, Transforms.elements_equal(shape, kept, equal)), input, plan)))
649
650# d(a slice of a)/da: g read with the kept shape, the slice's output with the
651# axis kept (PyTorch's unsqueeze for select), then scattered into +0.0 words
652# of a's shape (PyTorch's slice_backward): each input cell the slice read
653# gets the gradient word of the output cell it went to.
654def vjp_slice(+a: Nat, +input: Shape.Shape, +kept: Shape.Shape, +plan: Slice.Plan, g: Core.Dynamic) -> Pass:
655  Core.Packed{+shape, tensor} = g
656  slice_unsqueezed(a, input, kept, plan, shape, tensor, Shape.bool_case(Nat.is_eq(Shape.elements(shape), Shape.elements(kept))))
657
658def gather_sent(+a: Nat, +input: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>) -> Pass:
659  (g, result) = pair
660  match result:
661    case Fail{error}:
662      Pass{Unsaved{}, g, Fail{GatherError{error}}}
663    case Done{scattered}:
664      Pass{Unsaved{}, g, Done{[Contribution{a, input, scattered}]}}
665
666# d(an index_select or a gather of a)/da: g, of the output's shape, added
667# into +0.0 words of a's shape at the entries (PyTorch's index_add and
668# scatter_add): each input cell gets the sum of the gradient words of the
669# output cells that read it, in output order.
670def vjp_gather(+a: Nat, +input: Shape.Shape, +plan: Gather.Plan, +entries: List<&2, U32>, g: Core.Dynamic) -> Pass:
671  gather_sent(a, input, Gather.scatter(g, input, plan, entries))
672
673# The vector-Jacobian product of one node with gradient g: the contribution
674# each operand receives, in operand order, already reduced to its shape. A
675# constant sends nothing, and Frozen{id, node} sends node's contributions
676# without the ones to id.
677def vjp(node: Node, saved: Saved, g: Core.Dynamic) -> Pass:
678  match node:
679    case Leaf{}:
680      Pass{saved, g, Done{Nil{}}}
681    case Add{+a, +b, +left, +right}:
682      vjp_add(a, b, left, right, g)
683    case Sub{+a, +b, +left, +right}:
684      vjp_sub(a, b, left, right, g)
685    case Mul{+a, +b, +left, +right}:
686      vjp_mul(a, b, left, right, saved, g)
687    case Sum{+a, +input}:
688      vjp_sum(a, input, g)
689    case SumDims{+a, +input, +kept}:
690      vjp_sum_dims(a, input, kept, g)
691    case Matmul{+a, +b, +left, +right}:
692      vjp_matmul(a, b, left, right, saved, g)
693    case MatmulNT{+a, +b, +left, +right}:
694      vjp_matmul_nt(a, b, left, right, saved, g)
695    case Relu{+a, +input}:
696      vjp_relu(a, input, saved, g)
697    case MseLoss{+a, +b, +left, +right}:
698      vjp_mse(a, b, left, right, saved, g)
699    case CrossEntropy{+a, +input}:
700      vjp_cross_entropy(a, input, saved, g)
701    case BatchedMatmul{+a, +b, +left, +right}:
702      vjp_batched(a, b, left, right, saved, g)
703    case Slice{+a, +input, +kept, +plan}:
704      vjp_slice(a, input, kept, plan, g)
705    case Gathered{+a, +input, +plan, +entries}:
706      vjp_gather(a, input, plan, entries, g)
707    case Constant{}:
708      Pass{saved, g, Done{Nil{}}}
709    case Frozen{+id, inner}:
710      match inner:
711        case MatmulNT{+a, +b, +left, +right}:
712          nt_frozen(id, a, b, left, right, saved, g, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))))
713        case other:
714          unsent(id, vjp(other, saved, g))
715
716# Delivery.
717
718def placed(rest: List<&1, Gradient.Slot>, outcome: Gradient.Outcome) -> Result<&2, &1, Error, List<&1, Gradient.Slot>>:
719  match outcome:
720    case Gradient.Accumulated{slot}:
721      Done{slot <> rest}
722    case Gradient.Refused{slot, contribution, error}:
723      Fail{GradientError{error}}
724
725# Accumulates a contribution into the slot of node target, in a newest-first
726# list of count slots; same says whether the head is that node. A target the
727# list does not hold changes nothing.
728def bump(slots: List<&1, Gradient.Slot>, +count: Nat, +target: Nat, +shape: Shape.Shape, gradient: Core.Dynamic, same: Bool) ->
729  Result<&2, &1, Error, List<&1, Gradient.Slot>>:
730  match slots count same:
731    case Nil{} count same:
732      Done{Nil{}}
733    case Con{slot, rest} 0n same:
734      Done{slot <> rest}
735    case Con{slot, rest} 1n+below True{}:
736      placed(rest, Gradient.accumulate(shape, slot, gradient))
737    case Con{slot, rest} 1n+below False{}:
738      prepend(Gradient.Slot, slot, bump(rest, below, target, shape, gradient, at_head(below, target)))
739
740def deliver(contributions: List<&1, Contribution>, +count: Nat, state: Result<&2, &1, Error, List<&1, Gradient.Slot>>) ->
741  Result<&2, &1, Error, List<&1, Gradient.Slot>>:
742  match contributions state:
743    case Nil{} state:
744      state
745    case Con{contribution, rest} Fail{error}:
746      Fail{error}
747    case Con{Contribution{+target, +shape, gradient}, rest} Done{slots}:
748      deliver(rest, count, bump(slots, count, target, shape, gradient, at_head(count, target)))
749
750def sent(contributions: Result<&2, &1, Error, List<&1, Contribution>>, +count: Nat, slots: List<&1, Gradient.Slot>) ->
751  Result<&2, &1, Error, List<&1, Gradient.Slot>>:
752  match contributions:
753    case Fail{error}:
754      Fail{error}
755    case Done{items}:
756      deliver(items, count, Done{slots})
757
758# The sweep.
759
760def Swept() -> Type:
761  List<&1, Entry> & Result<&2, &1, Error, List<&1, Gradient.Slot>>
762
763def finished(entry: Entry, slot: Gradient.Slot, pair: Swept()) -> Swept():
764  (rest, result) = pair
765  (entry <> rest, prepend(Gradient.Slot, slot, result))
766
767def passed(node: Node, +shape: Shape.Shape, value: Core.Dynamic, slot: Gradient.Slot, +below: Nat,
768  pending: List<&1, Gradient.Slot>, rec: Result<&2, &1, Error, List<&1, Gradient.Slot>> -> Swept(), pass: Pass) -> Swept():
769  Pass{saved, g, contributions} = pass
770  finished(Entry{node, shape, value, saved, slot}, Gradient.Held{g}, rec(sent(contributions, below, pending)))
771
772def visited(+node: Node, +shape: Shape.Shape, value: Core.Dynamic, saved: Saved, slot: Gradient.Slot, gradient: Gradient.Slot,
773  +below: Nat, pending: List<&1, Gradient.Slot>, rec: Result<&2, &1, Error, List<&1, Gradient.Slot>> -> Swept()) -> Swept():
774  match gradient:
775    case Gradient.Unset{}:
776      finished(Entry{node, shape, value, saved, slot}, Gradient.Unset{}, rec(Done{pending}))
777    case Gradient.Held{g}:
778      passed(node, shape, value, slot, below, pending, rec, vjp(node, saved, g))
779
780# Visits entries newest first with their working gradients. Once a step
781# fails, the remaining entries pass through unchanged.
782def sweep(entries: List<&1, Entry>, state: Result<&2, &1, Error, List<&1, Gradient.Slot>>, +count: Nat) -> Swept():
783  match entries state count:
784    case Nil{} state count:
785      (Nil{}, state)
786    case Con{entry, rest} Fail{error} count:
787      (entry <> rest, Fail{error})
788    case Con{entry, rest} Done{Nil{}} count:
789      (entry <> rest, Done{Nil{}})
790    case Con{entry, rest} Done{Con{gradient, pending}} 0n:
791      (entry <> rest, Done{Nil{}})
792    case Con{Entry{+node, +shape, value, saved, slot}, rest} Done{Con{gradient, pending}} 1n+below:
793      visited(node, shape, value, saved, slot, gradient, below, pending, state => sweep(rest, state, below))
794
795def seeded(node: Node, +shape: Shape.Shape, value: Core.Dynamic, saved: Saved, slot: Gradient.Slot, rest: List<&1, Entry>,
796  +below: Nat, ones: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Swept():
797  match ones:
798    case Fail{error}:
799      (Entry{node, shape, value, saved, slot} <> rest, Fail{TensorError{error}})
800    case Done{tensor}:
801      sweep(Entry{node, shape, value, saved, slot} <> rest,
802        Done{Gradient.Held{Core.Packed{shape, tensor}} <> unsets(below)}, 1n+below)
803
804# The seed: ones of the output's shape.
805def seeded_entry(entry: Entry, rest: List<&1, Entry>, +below: Nat) -> Swept():
806  Entry{node, +shape, value, saved, slot} = entry
807  seeded(node, shape, value, saved, slot, rest, below, Core.new(shape, 1.0))
808
809# The working gradient of every entry, newest first: Unset for the entries
810# newer than output and for every entry output does not reach, ones of the
811# output's shape at output, and the sweep's totals below it.
812def pass(entries: List<&1, Entry>, +count: Nat, +output: Nat, selected: Bool) -> Swept():
813  match entries count selected:
814    case Nil{} count selected:
815      (Nil{}, Done{Nil{}})
816    case Con{entry, rest} 0n selected:
817      (entry <> rest, Done{Nil{}})
818    case Con{entry, rest} 1n+below False{}:
819      finished(entry, Gradient.Unset{}, pass(rest, below, output, at_head(below, output)))
820    case Con{entry, rest} 1n+below True{}:
821      seeded_entry(entry, rest, below)
822
823def operands_below(node: Node, +bound: Nat) -> Bool:
824  match node:
825    case Leaf{}:
826      True{}
827    case Add{a, b, left, right}:
828      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
829    case Sub{a, b, left, right}:
830      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
831    case Mul{a, b, left, right}:
832      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
833    case Sum{a, input}:
834      Nat.is_lt(a, bound)
835    case SumDims{a, input, kept}:
836      Nat.is_lt(a, bound)
837    case Matmul{a, b, left, right}:
838      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
839    case MatmulNT{a, b, left, right}:
840      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
841    case Relu{a, input}:
842      Nat.is_lt(a, bound)
843    case MseLoss{a, b, left, right}:
844      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
845    case CrossEntropy{a, input}:
846      Nat.is_lt(a, bound)
847    case BatchedMatmul{a, b, left, right}:
848      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
849    case Slice{a, input, kept, plan}:
850      Nat.is_lt(a, bound)
851    case Gathered{a, input, plan, entries}:
852      Nat.is_lt(a, bound)
853    case Constant{}:
854      True{}
855    case Frozen{id, inner}:
856      operands_below(inner, bound)
857
858def wired(nodes: List<&2, Node>, +count: Nat) -> Bool:
859  match nodes count:
860    case Nil{} count:
861      True{}
862    case Con{node, rest} 0n:
863      False{}
864    case Con{node, rest} 1n+below:
865      Bool.and(operands_below(node, below), wired(rest, below))
866
867# Accumulates each working gradient into its entry's slot. A refusal keeps
868# the slot and is reported after the walk.
869def merged_entry(node: Node, +shape: Shape.Shape, value: Core.Dynamic, saved: Saved, outcome: Gradient.Outcome,
870  pair: List<&1, Entry> & Result<&2, &2, Error, Unit>) -> List<&1, Entry> & Result<&2, &2, Error, Unit>:
871  match outcome:
872    case Gradient.Accumulated{slot}:
873      (rest, result) = pair
874      (Entry{node, shape, value, saved, slot} <> rest, result)
875    case Gradient.Refused{slot, contribution, error}:
876      (rest, result) = pair
877      (Entry{node, shape, value, saved, slot} <> rest, Fail{GradientError{error}})
878
879def merge(entries: List<&1, Entry>, gradients: List<&1, Gradient.Slot>) -> List<&1, Entry> & Result<&2, &2, Error, Unit>:
880  match entries gradients:
881    case Nil{} gradients:
882      (Nil{}, Done{Unit{}})
883    case Con{entry, rest} Nil{}:
884      (entry <> rest, Done{Unit{}})
885    case Con{Entry{node, +shape, value, saved, slot}, rest} Con{Gradient.Unset{}, more}:
886      merged_entry(node, shape, value, saved, Gradient.Accumulated{slot}, merge(rest, more))
887    case Con{Entry{node, +shape, value, saved, slot}, rest} Con{Gradient.Held{g}, more}:
888      merged_entry(node, shape, value, saved, Gradient.accumulate(shape, slot, g), merge(rest, more))
889
890def merged_tape(pair: List<&1, Entry> & Result<&2, &2, Error, Unit>) -> Tape & Result<&2, &2, Error, Unit>:
891  (entries, result) = pair
892  (Tape{entries}, result)
893
894def ran(pair: Swept()) -> Tape & Result<&2, &2, Error, Unit>:
895  (entries, result) = pair
896  match result:
897    case Fail{error}:
898      (Tape{entries}, Fail{error})
899    case Done{gradients}:
900      merged_tape(merge(entries, gradients))
901
902def checked(entries: List<&1, Entry>, +count: Nat, +output: Nat, known: Bool, sound: Bool) -> Tape & Result<&2, &2, Error, Unit>:
903  match known sound:
904    case False{} sound:
905      (Tape{entries}, Fail{UnknownNode{}})
906    case True{} False{}:
907      (Tape{entries}, Fail{DanglingOperand{}})
908    case True{} True{}:
909      ran(pass(entries, count, output, at_head(count, output)))
910
911def planned(+output: Nat, pair: List<&1, Entry> & List<&2, Node>) -> Tape & Result<&2, &2, Error, Unit>:
912  (entries, +nodes) = pair
913  +count = List.length(&2, Node, nodes)
914  checked(entries, count, output, Nat.is_lt(output, count), wired(nodes, count))
915
916# Runs one backward of output and accumulates its gradients into the slots.
917# An output the tape does not hold is refused with UnknownNode; a tape with a
918# node that refers to itself or to a newer or missing node is refused with
919# DanglingOperand. Both refusals leave the tape unchanged.
920def backward(tape: Tape, +output: Nat) -> Tape & Result<&2, &2, Error, Unit>:
921  Tape{entries} = tape
922  planned(output, listed(entries))
923
924def cleared(entries: List<&1, Entry>) -> List<&1, Entry>:
925  match entries:
926    case Nil{}:
927      Nil{}
928    case Con{Entry{node, shape, value, saved, slot}, rest}:
929      Entry{node, shape, value, saved, Gradient.reset(slot)} <> cleared(rest)
930
931# Empties every gradient slot, like PyTorch's zero_grad(set_to_none=True).
932def zero_grad(tape: Tape) -> Tape:
933  Tape{entries} = tape
934  Tape{cleared(entries)}
935
936# Reads.
937
938def Found() -> Type:
939  List<&1, Entry> & Result<&2, &1, Error, Core.Dynamic>
940
941def cloned_value(node: Node, +shape: Shape.Shape, saved: Saved, slot: Gradient.Slot, rest: List<&1, Entry>,
942  pair: Core.Dynamic & Core.Dynamic) -> Found():
943  (value, copy) = pair
944  (Entry{node, shape, value, saved, slot} <> rest, Done{copy})
945
946def found_value(entry: Entry, rest: List<&1, Entry>) -> Found():
947  Entry{node, +shape, value, saved, slot} = entry
948  cloned_value(node, shape, saved, slot, rest, Core.clone_dynamic(value))
949
950def prepend_found(entry: Entry, pair: Found()) -> Found():
951  (rest, result) = pair
952  (entry <> rest, result)
953
954# A copy of the value of node id, in a newest-first list of count entries.
955def value_at(entries: List<&1, Entry>, +count: Nat, +id: Nat, same: Bool) -> Found():
956  match entries count same:
957    case Nil{} count same:
958      (Nil{}, Fail{UnknownNode{}})
959    case Con{entry, rest} 0n same:
960      (entry <> rest, Fail{UnknownNode{}})
961    case Con{entry, rest} 1n+below True{}:
962      found_value(entry, rest)
963    case Con{entry, rest} 1n+below False{}:
964      prepend_found(entry, value_at(rest, below, id, at_head(below, id)))
965
966def counted_entry(entry: Entry, pair: List<&1, Entry> & Nat) -> List<&1, Entry> & Nat:
967  (rest, count) = pair
968  (entry <> rest, 1n+count)
969
970# The number of entries, next to the entries.
971def counted(entries: List<&1, Entry>) -> List<&1, Entry> & Nat:
972  match entries:
973    case Nil{}:
974      (Nil{}, 0n)
975    case Con{entry, rest}:
976      counted_entry(entry, counted(rest))
977
978def value_counted(+id: Nat, pair: List<&1, Entry> & Nat) -> Found():
979  (entries, +count) = pair
980  value_at(entries, count, id, at_head(count, id))
981
982def found_tape(pair: Found()) -> Tape & Result<&2, &1, Error, Core.Dynamic>:
983  (entries, result) = pair
984  (Tape{entries}, result)
985
986# A copy of the forward value of node id.
987def value(tape: Tape, +id: Nat) -> Tape & Result<&2, &1, Error, Core.Dynamic>:
988  Tape{entries} = tape
989  found_tape(value_counted(id, counted(entries)))
990
991def Graded() -> Type:
992  List<&1, Entry> & Result<&2, &1, Error, Maybe<&1, Core.Dynamic>>
993
994def held_copy(pair: Core.Dynamic & Core.Dynamic) -> Gradient.Slot & Maybe<&1, Core.Dynamic>:
995  (gradient, copy) = pair
996  (Gradient.Held{gradient}, Some{copy})
997
998def copied_grad(read: Maybe<&1, Core.Dynamic>) -> Gradient.Slot & Maybe<&1, Core.Dynamic>:
999  match read:
1000    case None{}:
1001      (Gradient.Unset{}, None{})
1002    case Some{gradient}:
1003      held_copy(Core.clone_dynamic(gradient))
1004
1005def graded_entry(node: Node, +shape: Shape.Shape, value: Core.Dynamic, saved: Saved, rest: List<&1, Entry>,
1006  pair: Gradient.Slot & Maybe<&1, Core.Dynamic>) -> Graded():
1007  (slot, copy) = pair
1008  (Entry{node, shape, value, saved, slot} <> rest, Done{copy})
1009
1010def found_grad(entry: Entry, rest: List<&1, Entry>) -> Graded():
1011  Entry{node, +shape, value, saved, slot} = entry
1012  graded_entry(node, shape, value, saved, rest, copied_grad(Gradient.grad(slot)))
1013
1014def prepend_graded(entry: Entry, pair: Graded()) -> Graded():
1015  (rest, result) = pair
1016  (entry <> rest, result)
1017
1018def grad_at(entries: List<&1, Entry>, +count: Nat, +id: Nat, same: Bool) -> Graded():
1019  match entries count same:
1020    case Nil{} count same:
1021      (Nil{}, Fail{UnknownNode{}})
1022    case Con{entry, rest} 0n same:
1023      (entry <> rest, Fail{UnknownNode{}})
1024    case Con{entry, rest} 1n+below True{}:
1025      found_grad(entry, rest)
1026    case Con{entry, rest} 1n+below False{}:
1027      prepend_graded(entry, grad_at(rest, below, id, at_head(below, id)))
1028
1029def grad_counted(+id: Nat, pair: List<&1, Entry> & Nat) -> Graded():
1030  (entries, +count) = pair
1031  grad_at(entries, count, id, at_head(count, id))
1032
1033def graded_tape(pair: Graded()) -> Tape & Result<&2, &1, Error, Maybe<&1, Core.Dynamic>>:
1034  (entries, result) = pair
1035  (Tape{entries}, result)
1036
1037# A copy of the gradient slot of node id: None until a backward reaches it.
1038def grad(tape: Tape, +id: Nat) -> Tape & Result<&2, &1, Error, Maybe<&1, Core.Dynamic>>:
1039  Tape{entries} = tape
1040  graded_tape(grad_counted(id, counted(entries)))
1041
1042# Recording.
1043
1044def pushed(node: Node, value: Core.Dynamic, saved: Saved, pair: List<&1, Entry> & Nat) -> Tape & Nat:
1045  Core.Packed{+shape, tensor} = value
1046  (entries, +count) = pair
1047  (Tape{Entry{node, shape, Core.Packed{shape, tensor}, saved, Gradient.Unset{}} <> entries}, count)
1048
1049def done_id(pair: Tape & Nat) -> Tape & Result<&2, &2, Error, Nat>:
1050  (tape, id) = pair
1051  (tape, Done{id})
1052
1053def is_constant(node: Node) -> Bool:
1054  match node:
1055    case Constant{}:
1056      True{}
1057    case other:
1058      False{}
1059
1060# Whether node id of a newest-first list of count nodes is a constant.
1061def constant_at(nodes: List<&2, Node>, +count: Nat, +id: Nat, same: Bool) -> Bool:
1062  match nodes count same:
1063    case Nil{} count same:
1064      False{}
1065    case Con{node, rest} 0n same:
1066      False{}
1067    case Con{node, rest} 1n+below True{}:
1068      is_constant(node)
1069    case Con{node, rest} 1n+below False{}:
1070      constant_at(rest, below, id, at_head(below, id))
1071
1072# The operand ids of a node, in operand order.
1073def operand_ids(node: Node) -> List<&2, Nat>:
1074  match node:
1075    case Leaf{}:
1076      []
1077    case Add{a, b, left, right}:
1078      [a, b]
1079    case Sub{a, b, left, right}:
1080      [a, b]
1081    case Mul{a, b, left, right}:
1082      [a, b]
1083    case Sum{a, input}:
1084      [a]
1085    case SumDims{a, input, kept}:
1086      [a]
1087    case Matmul{a, b, left, right}:
1088      [a, b]
1089    case MatmulNT{a, b, left, right}:
1090      [a, b]
1091    case Relu{a, input}:
1092      [a]
1093    case MseLoss{a, b, left, right}:
1094      [a, b]
1095    case CrossEntropy{a, input}:
1096      [a]
1097    case BatchedMatmul{a, b, left, right}:
1098      [a, b]
1099    case Slice{a, input, kept, plan}:
1100      [a]
1101    case Gathered{a, input, plan, entries}:
1102      [a]
1103    case Constant{}:
1104      []
1105    case Frozen{id, inner}:
1106      operand_ids(inner)
1107
1108def frozen_if(+id: Nat, node: Node, constant: Bool) -> Node:
1109  match constant:
1110    case True{}:
1111      Frozen{id, node}
1112    case False{}:
1113      node
1114
1115# node wrapped in Frozen{id, .} for each operand id that is a constant of the
1116# newest-first list of count nodes, the first operand outermost.
1117def freeze(ids: List<&2, Nat>, +nodes: List<&2, Node>, +count: Nat, node: Node) -> Node:
1118  match ids:
1119    case Nil{}:
1120      node
1121    case Con{+id, rest}:
1122      frozen_if(id, freeze(rest, nodes, count, node), constant_at(nodes, count, id, at_head(count, id)))
1123
1124# Whether every id of a list is a constant of the newest-first list of count
1125# nodes.
1126def all_constant(ids: List<&2, Nat>, +nodes: List<&2, Node>, +count: Nat) -> Bool:
1127  match ids:
1128    case Nil{}:
1129      True{}
1130    case Con{+id, rest}:
1131      Bool.and(constant_at(nodes, count, id, at_head(count, id)), all_constant(rest, nodes, count))
1132
1133# A node whose operands are all constants is a constant itself, as in PyTorch
1134# an operation on tensors that do not require grad gives one that does not;
1135# any other node is frozen at each operand that is a constant.
1136def settled(+node: Node, +ids: List<&2, Nat>, +nodes: List<&2, Node>, +count: Nat, constant: Bool) -> Node:
1137  match constant:
1138    case True{}:
1139      Constant{}
1140    case False{}:
1141      freeze(ids, nodes, count, node)
1142
1143def marked(+node: Node, value: Core.Dynamic, saved: Saved, pair: List<&1, Entry> & List<&2, Node>) -> Tape & Nat:
1144  (entries, +nodes) = pair
1145  +count = List.length(&2, Node, nodes)
1146  +ids = operand_ids(node)
1147  pushed(settled(node, ids, nodes, count, all_constant(ids, nodes, count)), value, saved, (entries, count))
1148
1149# The new entry's id is the number of entries before it. A recording has at
1150# least one operand; when they are all constants it records a constant, and
1151# otherwise each operand that is a constant freezes the node.
1152def push(node: Node, value: Core.Dynamic, saved: Saved, entries: List<&1, Entry>) -> Tape & Result<&2, &2, Error, Nat>:
1153  done_id(marked(node, value, saved, listed(entries)))
1154
1155# Records a leaf.
1156def leaf(tape: Tape, value: Core.Dynamic) -> Tape & Nat:
1157  Tape{entries} = tape
1158  pushed(Leaf{}, value, Unsaved{}, counted(entries))
1159
1160# Records a constant: a leaf that takes no gradient, PyTorch's tensor with
1161# requires_grad=False. No VJP sends it a contribution, so its slot reads None
1162# after any backward that does not start at it; Sgd.step takes it as a leaf
1163# and keeps it.
1164def constant(tape: Tape, value: Core.Dynamic) -> Tape & Nat:
1165  Tape{entries} = tape
1166  pushed(Constant{}, value, Unsaved{}, counted(entries))
1167
1168def recorded(node: Node, saved: Saved, entries: List<&1, Entry>, result: Result<&2, &1, Core.Error, Core.Dynamic>) ->
1169  Tape & Result<&2, &2, Error, Nat>:
1170  match result:
1171    case Fail{error}:
1172      (Tape{entries}, Fail{TensorError{error}})
1173    case Done{output}:
1174      push(node, output, saved, entries)
1175
1176def multiplied_record(node: Node, entries: List<&1, Entry>,
1177  result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Tape & Result<&2, &2, Error, Nat>:
1178  match result:
1179    case Fail{error}:
1180      (Tape{entries}, Fail{MatrixError{error}})
1181    case Done{triple}:
1182      (left, right, output) = triple
1183      push(node, output, Operands{left, right}, entries)
1184
1185def dropped(node: Node, entries: List<&1, Entry>, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
1186  Tape & Result<&2, &2, Error, Nat>:
1187  (operands, result) = pair
1188  recorded(node, Unsaved{}, entries, result)
1189
1190def kept(node: Node, entries: List<&1, Entry>, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) ->
1191  Tape & Result<&2, &2, Error, Nat>:
1192  (operands, result) = pair
1193  (left, right) = operands
1194  recorded(node, Operands{left, right}, entries, result)
1195
1196def loss_recorded(node: Node, entries: List<&1, Entry>, result: Result<&2, &1, Loss.Error, Loss.Recorded>) ->
1197  Tape & Result<&2, &2, Error, Nat>:
1198  match result:
1199    case Fail{error}:
1200      (Tape{entries}, Fail{LossError{error}})
1201    case Done{Loss.Recorded{loss, left, right}}:
1202      push(node, loss, Operands{left, right}, entries)
1203
1204def loss_dropped(node: Node, entries: List<&1, Entry>, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Loss.Error, Loss.Recorded>) ->
1205  Tape & Result<&2, &2, Error, Nat>:
1206  (operands, result) = pair
1207  loss_recorded(node, entries, result)
1208
1209def loss_kept(node: Node, entries: List<&1, Entry>, pair: Core.Dynamic & Result<&2, &1, Loss.Error, Loss.Recorded>) ->
1210  Tape & Result<&2, &2, Error, Nat>:
1211  (input, result) = pair
1212  loss_recorded(node, entries, result)
1213
1214# torch.matmul on the tape: two rank-2 operands record `Matmul`, the
1215# product `matmul` records, and every other pair records `BatchedMatmul`.
1216def batched_record(+a: Nat, +b: Nat, +left_shape: Shape.Shape, +right_shape: Shape.Shape, entries: List<&1, Entry>,
1217  left: Core.Dynamic, right: Core.Dynamic, matrices: Bool) -> Tape & Result<&2, &2, Error, Nat>:
1218  match matrices:
1219    case True{}:
1220      multiplied_record(Matmul{a, b, left_shape, right_shape}, entries, Matrix.matmul(left, right))
1221    case False{}:
1222      multiplied_record(BatchedMatmul{a, b, left_shape, right_shape}, entries, Batched.matmul(left, right))
1223
1224def operated(kind: Binary, +a: Nat, +b: Nat, +left_shape: Shape.Shape, +right_shape: Shape.Shape, entries: List<&1, Entry>,
1225  left: Core.Dynamic, right: Core.Dynamic) -> Tape & Result<&2, &2, Error, Nat>:
1226  match kind:
1227    case Plus{}:
1228      dropped(Add{a, b, left_shape, right_shape}, entries, Elementwise.add(left, right))
1229    case Minus{}:
1230      dropped(Sub{a, b, left_shape, right_shape}, entries, Elementwise.sub(left, right))
1231    case Times{}:
1232      kept(Mul{a, b, left_shape, right_shape}, entries, Elementwise.mul(left, right))
1233    case Product{}:
1234      multiplied_record(Matmul{a, b, left_shape, right_shape}, entries, Matrix.matmul(left, right))
1235    case ProductNT{}:
1236      multiplied_record(MatmulNT{a, b, left_shape, right_shape}, entries, Matrix.matmul_nt(left, right))
1237    case SquaredError{}:
1238      loss_dropped(MseLoss{a, b, left_shape, right_shape}, entries, Loss.mse_loss(left, right))
1239    case BatchedProduct{}:
1240      batched_record(a, b, left_shape, right_shape, entries, left, right,
1241        Batched.both_matrices(left_shape, right_shape))
1242
1243def computed(kind: Binary, +a: Nat, +b: Nat, entries: List<&1, Entry>, left: Core.Dynamic & Shape.Shape,
1244  right: Core.Dynamic & Shape.Shape) -> Tape & Result<&2, &2, Error, Nat>:
1245  (left, +left_shape) = left
1246  (right, +right_shape) = right
1247  operated(kind, a, b, left_shape, right_shape, entries, left, right)
1248
1249def right_found(kind: Binary, +a: Nat, +b: Nat, left: Core.Dynamic, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1250  (entries, right) = pair
1251  match right:
1252    case Fail{error}:
1253      (Tape{entries}, Fail{UnknownNode{}})
1254    case Done{right}:
1255      computed(kind, a, b, entries, Core.shape(left), Core.shape(right))
1256
1257def left_found(kind: Binary, +a: Nat, +b: Nat, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1258  (entries, left) = pair
1259  match left:
1260    case Fail{error}:
1261      (Tape{entries}, Fail{UnknownNode{}})
1262    case Done{left}:
1263      right_found(kind, a, b, left, value_counted(b, counted(entries)))
1264
1265# Records `a op b` from copies of the operands' values. An operand id the tape
1266# does not hold is refused with UnknownNode, and an operation the tensor layer
1267# refuses (shapes that do not broadcast, or a matmul of the wrong ranks or
1268# inner dimensions) with its error; the tape comes back unchanged.
1269def binary(tape: Tape, kind: Binary, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1270  Tape{entries} = tape
1271  left_found(kind, a, b, value_counted(a, counted(entries)))
1272
1273def add(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1274  binary(tape, Plus{}, a, b)
1275
1276def sub(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1277  binary(tape, Minus{}, a, b)
1278
1279def mul(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1280  binary(tape, Times{}, a, b)
1281
1282def matmul(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1283  binary(tape, Product{}, a, b)
1284
1285# Records torch.matmul(a, b): batched products and 1-D promotion
1286# (Batched.matmul). Two rank-2 operands record exactly what `matmul` records.
1287# A rank-0 operand is refused with RankMismatch (105), unequal inner extents
1288# with InnerDimensionMismatch (106), batch dimensions that do not broadcast
1289# with IncompatibleBroadcast (6); the tape comes back unchanged.
1290def batched_matmul(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1291  binary(tape, BatchedProduct{}, a, b)
1292
1293# Records a @ b^T, with b read in place: PyTorch's linear(a, b) without a
1294# bias.
1295def matmul_nt(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1296  binary(tape, ProductNT{}, a, b)
1297
1298# The tape without its newest entry.
1299def popped(tape: Tape) -> Tape:
1300  Tape{entries} = tape
1301  match entries:
1302    case Nil{}:
1303      Tape{Nil{}}
1304    case Con{entry, rest}:
1305      Tape{rest}
1306
1307# A refused sum drops the product, the newest entry, so the tape comes back
1308# as it was before the product was recorded.
1309def unrecorded(pair: Tape & Result<&2, &2, Error, Nat>) -> Tape & Result<&2, &2, Error, Nat>:
1310  (tape, result) = pair
1311  match result:
1312    case Fail{error}:
1313      (popped(tape), Fail{error})
1314    case Done{id}:
1315      (tape, Done{id})
1316
1317def biased(+bias: Nat, pair: Tape & Result<&2, &2, Error, Nat>) -> Tape & Result<&2, &2, Error, Nat>:
1318  (tape, result) = pair
1319  match result:
1320    case Fail{error}:
1321      (tape, Fail{error})
1322    case Done{+product}:
1323      unrecorded(add(tape, product, bias))
1324
1325def linear_checked(+x: Nat, +w: Nat, +bias: Nat, entries: List<&1, Entry>, known: Bool) -> Tape & Result<&2, &2, Error, Nat>:
1326  match known:
1327    case False{}:
1328      (Tape{entries}, Fail{UnknownNode{}})
1329    case True{}:
1330      biased(bias, matmul_nt(Tape{entries}, x, w))
1331
1332def linear_counted(+x: Nat, +w: Nat, +bias: Nat, pair: List<&1, Entry> & Nat) -> Tape & Result<&2, &2, Error, Nat>:
1333  (entries, +count) = pair
1334  linear_checked(x, w, bias, entries, Nat.is_lt(bias, count))
1335
1336# Linear: records x @ W^T with W of [out, in], PyTorch's layout, then adds
1337# the bias b of [out] to every row. The id is the sum's; the product is the
1338# entry before it. The bias id is checked against the tape the call receives,
1339# before the product is recorded, so the product's own id never stands in
1340# for a missing bias. A refusal of any step (an unknown id, a product or a
1341# sum the tensor layer refuses) comes back with its error and the tape as it
1342# was.
1343def linear(tape: Tape, +x: Nat, +w: Nat, +bias: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1344  Tape{entries} = tape
1345  linear_counted(x, w, bias, counted(entries))
1346
1347def summed(+a: Nat, entries: List<&1, Entry>, +shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>) ->
1348  Tape & Result<&2, &2, Error, Nat>:
1349  match result:
1350    case Fail{error}:
1351      (Tape{entries}, Fail{TensorError{error}})
1352    case Done{pair}:
1353      (input, output) = pair
1354      push(Sum{a, shape}, output, Unsaved{}, entries)
1355
1356def sum_shaped(+a: Nat, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) -> Tape & Result<&2, &2, Error, Nat>:
1357  (input, +shape) = pair
1358  summed(a, entries, shape, Reduction.sum_dynamic(input))
1359
1360def sum_found(+a: Nat, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1361  (entries, input) = pair
1362  match input:
1363    case Fail{error}:
1364      (Tape{entries}, Fail{UnknownNode{}})
1365    case Done{input}:
1366      sum_shaped(a, entries, Core.shape(input))
1367
1368# Records the sum of every element of a, a rank-0 tensor.
1369def sum(tape: Tape, +a: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1370  Tape{entries} = tape
1371  sum_found(a, value_counted(a, counted(entries)))
1372
1373def dims_error(error: SumDims.Error) -> Error:
1374  match error:
1375    case SumDims.ReduceError{error}:
1376      ReduceError{error}
1377    case SumDims.TransformError{error}:
1378      TransformError{error}
1379
1380def dims_summed(node: Node, entries: List<&1, Entry>, pair: Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>) ->
1381  Tape & Result<&2, &2, Error, Nat>:
1382  (input, result) = pair
1383  match result:
1384    case Fail{error}:
1385      (Tape{entries}, Fail{dims_error(error)})
1386    case Done{output}:
1387      push(node, output, Unsaved{}, entries)
1388
1389def dims_planned(+a: Nat, entries: List<&1, Entry>, input: Core.Dynamic, +shape: Shape.Shape, keepdim: Bool,
1390  planned: Result<&2, &2, SumDims.Error, SumDims.Plan>) -> Tape & Result<&2, &2, Error, Nat>:
1391  match planned:
1392    case Fail{error}:
1393      (Tape{entries}, Fail{dims_error(error)})
1394    case Done{SumDims.Plan{+kept, +dropped}}:
1395      dims_summed(SumDims{a, shape, kept}, entries, SumDims.sum_planned(input, SumDims.Plan{kept, dropped}, keepdim))
1396
1397# The signed axes, each wrapped among the operand's axes (its rank, or 1 at
1398# rank 0), as PyTorch wraps them, before the plan checks them.
1399def dims_shaped(+a: Nat, +axes: List<&2, Signed.Position>, keepdim: Bool, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) ->
1400  Tape & Result<&2, &2, Error, Nat>:
1401  (input, +shape) = pair
1402  dims_planned(a, entries, input, shape, keepdim, SumDims.plan(shape, Signed.wrap_all(SumDims.limit(Shape.dimensions(shape)), axes)))
1403
1404def dims_found(+a: Nat, +axes: List<&2, Signed.Position>, keepdim: Bool, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1405  (entries, input) = pair
1406  match input:
1407    case Fail{error}:
1408      (Tape{entries}, Fail{UnknownNode{}})
1409    case Done{input}:
1410      dims_shaped(a, axes, keepdim, entries, Core.shape(input))
1411
1412# Records PyTorch's a.sum(dim=axes, keepdim=keepdim), SumDims.sum_dims of a
1413# copy of a's value: the kept shape when keepdim is True, else the dropped
1414# one. The axes are PyTorch's signed ones: Last{0} is -1, a's last axis. An
1415# empty axis list sums every axis. An axis outside -limit..limit-1, where
1416# limit is max(rank, 1), so a rank-0 a takes -1 and 0, is refused with
1417# AxisOutOfBounds (13), two axes that name one axis with DuplicateAxis (14),
1418# an unknown id with UnknownNode; the tape comes back unchanged.
1419def sum_dims(tape: Tape, +a: Nat, +axes: List<&2, Signed.Position>, keepdim: Bool) -> Tape & Result<&2, &2, Error, Nat>:
1420  Tape{entries} = tape
1421  dims_found(a, axes, keepdim, value_counted(a, counted(entries)))
1422
1423def slice_kept(node: Node, entries: List<&1, Entry>, result: Result<&2, &1, Slice.Error, Core.Dynamic>) ->
1424  Tape & Result<&2, &2, Error, Nat>:
1425  match result:
1426    case Fail{error}:
1427      (Tape{entries}, Fail{SliceError{error}})
1428    case Done{output}:
1429      push(node, output, Unsaved{}, entries)
1430
1431def slice_planned(+a: Nat, +shape: Shape.Shape, entries: List<&1, Entry>, result: Result<&2, &1, Slice.Error, Core.Dynamic>,
1432  planned: Result<&2, &2, Slice.Error, Slice.Planned>) -> Tape & Result<&2, &2, Error, Nat>:
1433  match planned:
1434    case Fail{error}:
1435      (Tape{entries}, Fail{SliceError{error}})
1436    case Done{Slice.Planned{+plan, +kept}}:
1437      slice_kept(Slice{a, shape, kept, plan}, entries, result)
1438
1439def slice_taken(+a: Nat, +shape: Shape.Shape, entries: List<&1, Entry>, planned: Result<&2, &2, Slice.Error, Slice.Planned>,
1440  pair: Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>) -> Tape & Result<&2, &2, Error, Nat>:
1441  (input, result) = pair
1442  slice_planned(a, shape, entries, result, planned)
1443
1444# Which of the three a slice step is: PyTorch's a[..., start:stop:step, ...]
1445# at axis `dim`, a.narrow(dim, start, length) and a.select(dim, index). The
1446# axis and the positions along it are PyTorch's signed ints (signed.bend).
1447type Slicing is Data:
1448  Sliced{dim: Signed.Position, start: Signed.Position, stop: Signed.Position, step: U32}
1449  Narrowed{dim: Signed.Position, start: Signed.Position, length: U32}
1450  Selected{dim: Signed.Position, index: Signed.Position}
1451
1452def slice_extent(+shape: Shape.Shape, +axis: U32) -> U32:
1453  Signed.extent(Shape.dimensions(shape), axis)
1454
1455def slice_axis_plan(+shape: Shape.Shape, +axis: U32, +slicing: Slicing) -> Result<&2, &2, Slice.Error, Slice.Planned>:
1456  match slicing:
1457    case Sliced{dim, start, stop, +step}:
1458      Slice.plan(shape, axis, Signed.bound(slice_extent(shape, axis), start), Signed.bound(slice_extent(shape, axis), stop), step)
1459    case Narrowed{dim, start, +length}:
1460      +first = Signed.wrap(slice_extent(shape, axis), start)
1461      Slice.plan(shape, axis, first, U32.add(first, length), 1)
1462    case Selected{dim, index}:
1463      +first = Signed.wrap(slice_extent(shape, axis), index)
1464      Slice.plan(shape, axis, first, U32.add(first, 1), 1)
1465
1466def slice_dim(+slicing: Slicing) -> Signed.Position:
1467  match slicing:
1468    case Sliced{dim, start, stop, step}:
1469      dim
1470    case Narrowed{dim, start, length}:
1471      dim
1472    case Selected{dim, index}:
1473      dim
1474
1475# The slice each one records: its plan, at the axis and positions the value
1476# takes (Signed.slice, narrow and select), and the value it computes. The
1477# value is checked first, so a narrow whose start lies before -extent, which
1478# only the value refuses, records nothing.
1479def slice_plan(+shape: Shape.Shape, +slicing: Slicing) -> Result<&2, &2, Slice.Error, Slice.Planned>:
1480  slice_axis_plan(shape, Signed.wrap(Shape.rank(shape), slice_dim(slicing)), slicing)
1481
1482def slice_value(input: Core.Dynamic, +slicing: Slicing) -> Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
1483  match slicing:
1484    case Sliced{+dim, +start, +stop, +step}:
1485      Signed.slice(input, dim, start, stop, step)
1486    case Narrowed{+dim, +start, +length}:
1487      Signed.narrow(input, dim, start, length)
1488    case Selected{+dim, +index}:
1489      Signed.select(input, dim, index)
1490
1491def slice_shaped(+a: Nat, +slicing: Slicing, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) -> Tape & Result<&2, &2, Error, Nat>:
1492  (input, +shape) = pair
1493  slice_taken(a, shape, entries, slice_plan(shape, slicing), slice_value(input, slicing))
1494
1495def slice_found(+a: Nat, +slicing: Slicing, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1496  (entries, input) = pair
1497  match input:
1498    case Fail{error}:
1499      (Tape{entries}, Fail{UnknownNode{}})
1500    case Done{input}:
1501      slice_shaped(a, slicing, entries, Core.shape(input))
1502
1503# Records a slice of a copy of a's value (Signed.slice, narrow or select) as
1504# a Slice node with a's shape, the slice's shape with the axis kept and its
1505# plan. The value is checked first, so a refusal is the value's: an axis
1506# outside -rank..rank-1 is AxisOutOfBounds (13), a step of 0 StepNotPositive
1507# (19), a narrow past the extent or with a start before -extent, or a select
1508# index outside -extent..extent-1, OutOfRange (20), an unknown id
1509# UnknownNode; the tape comes back unchanged.
1510def sliced(tape: Tape, +a: Nat, +slicing: Slicing) -> Tape & Result<&2, &2, Error, Nat>:
1511  Tape{entries} = tape
1512  slice_found(a, slicing, value_counted(a, counted(entries)))
1513
1514# Records a[..., start:stop:step, ...] with the slice at axis `dim`.
1515def slice(tape: Tape, +a: Nat, +dim: Signed.Position, +start: Signed.Position, +stop: Signed.Position, +step: U32) ->
1516  Tape & Result<&2, &2, Error, Nat>:
1517  sliced(tape, a, Sliced{dim, start, stop, step})
1518
1519# Records a.narrow(dim, start, length).
1520def narrow(tape: Tape, +a: Nat, +dim: Signed.Position, +start: Signed.Position, +length: U32) -> Tape & Result<&2, &2, Error, Nat>:
1521  sliced(tape, a, Narrowed{dim, start, length})
1522
1523# Records a.select(dim, index), which drops the axis.
1524def select(tape: Tape, +a: Nat, +dim: Signed.Position, +index: Signed.Position) -> Tape & Result<&2, &2, Error, Nat>:
1525  sliced(tape, a, Selected{dim, index})
1526
1527def gather_kept(node: Node, entries: List<&1, Entry>, result: Result<&2, &1, Gather.Error, Core.Dynamic>) ->
1528  Tape & Result<&2, &2, Error, Nat>:
1529  match result:
1530    case Fail{error}:
1531      (Tape{entries}, Fail{GatherError{error}})
1532    case Done{output}:
1533      push(node, output, Unsaved{}, entries)
1534
1535# Which of the two an indexing step is: PyTorch's a.index_select(dim, index)
1536# with a vector of entries, and a.gather(dim, index) with an index tensor.
1537type Indexing is Data:
1538  IndexSelected{dim: Signed.Position, entries: List<&2, U32>}
1539  GatheredBy{dim: Signed.Position, index: Gather.Index}
1540
1541def indexing_plan(+shape: Shape.Shape, +indexing: Indexing) -> Result<&2, &2, Gather.Error, Gather.Planned>:
1542  match indexing:
1543    case IndexSelected{+dim, +entries}:
1544      Gather.select_plan(shape, dim, entries)
1545    case GatheredBy{+dim, +index}:
1546      Gather.gather_plan(shape, dim, index)
1547
1548def indexing_value(input: Core.Dynamic, +indexing: Indexing) -> Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>:
1549  match indexing:
1550    case IndexSelected{+dim, +entries}:
1551      Gather.index_select(input, dim, entries)
1552    case GatheredBy{+dim, +index}:
1553      Gather.gather(input, dim, index)
1554
1555def indexing_entries(+indexing: Indexing) -> List<&2, U32>:
1556  match indexing:
1557    case IndexSelected{dim, entries}:
1558      entries
1559    case GatheredBy{dim, Gather.Index{shape, entries}}:
1560      entries
1561
1562def gather_planned(+a: Nat, +shape: Shape.Shape, +indexing: Indexing, entries: List<&1, Entry>,
1563  result: Result<&2, &1, Gather.Error, Core.Dynamic>, planned: Result<&2, &2, Gather.Error, Gather.Planned>) ->
1564  Tape & Result<&2, &2, Error, Nat>:
1565  match planned:
1566    case Fail{error}:
1567      (Tape{entries}, Fail{GatherError{error}})
1568    case Done{Gather.Planned{+plan, +output}}:
1569      gather_kept(Gathered{a, shape, plan, indexing_entries(indexing)}, entries, result)
1570
1571def gather_taken(+a: Nat, +shape: Shape.Shape, +indexing: Indexing, entries: List<&1, Entry>,
1572  pair: Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>) -> Tape & Result<&2, &2, Error, Nat>:
1573  (input, result) = pair
1574  gather_planned(a, shape, indexing, entries, result, indexing_plan(shape, indexing))
1575
1576def gather_shaped(+a: Nat, +indexing: Indexing, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) ->
1577  Tape & Result<&2, &2, Error, Nat>:
1578  (input, +shape) = pair
1579  gather_taken(a, shape, indexing, entries, indexing_value(input, indexing))
1580
1581def gather_found(+a: Nat, +indexing: Indexing, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1582  (entries, input) = pair
1583  match input:
1584    case Fail{error}:
1585      (Tape{entries}, Fail{UnknownNode{}})
1586    case Done{input}:
1587      gather_shaped(a, indexing, entries, Core.shape(input))
1588
1589# Records an index_select or a gather of a copy of a's value as a Gathered
1590# node with a's shape, the plan and the entries. The value is checked first,
1591# so a refusal is the value's: an axis outside -rank..rank-1 is
1592# AxisOutOfBounds (13), an entry at or past the extent OutOfRange (20), a
1593# gather index of another rank IndexRankMismatch (21), with another extent at
1594# another axis IndexShapeMismatch (22), whose entry count is not the product
1595# of its dimensions ValueCountMismatch (10), an unknown id UnknownNode; the
1596# tape comes back unchanged.
1597def indexed(tape: Tape, +a: Nat, +indexing: Indexing) -> Tape & Result<&2, &2, Error, Nat>:
1598  Tape{entries} = tape
1599  gather_found(a, indexing, value_counted(a, counted(entries)))
1600
1601# Records a.index_select(dim, index).
1602def index_select(tape: Tape, +a: Nat, +dim: Signed.Position, +entries: List<&2, U32>) -> Tape & Result<&2, &2, Error, Nat>:
1603  indexed(tape, a, IndexSelected{dim, entries})
1604
1605# Records a.gather(dim, index).
1606def gather(tape: Tape, +a: Nat, +dim: Signed.Position, +index: Gather.Index) -> Tape & Result<&2, &2, Error, Nat>:
1607  indexed(tape, a, GatheredBy{dim, index})
1608
1609def relu_zeroed(+a: Nat, entries: List<&1, Entry>, +shape: Shape.Shape, input: Core.Dynamic,
1610  zero: Result<&2, &1, Core.Error, Core.Tensor<Reduction.scalar_shape()>>) -> Tape & Result<&2, &2, Error, Nat>:
1611  match zero:
1612    case Fail{error}:
1613      (Tape{entries}, Fail{TensorError{error}})
1614    case Done{tensor}:
1615      kept(Relu{a, shape}, entries,
1616        Elementwise.binary(Elementwise.ReluLeft{}, input, Core.pack(Reduction.scalar_shape(), tensor)))
1617
1618def relu_shaped(+a: Nat, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) -> Tape & Result<&2, &2, Error, Nat>:
1619  (input, +shape) = pair
1620  relu_zeroed(a, entries, shape, input, Core.new(Reduction.scalar_shape(), 0.0))
1621
1622def relu_found(+a: Nat, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1623  (entries, input) = pair
1624  match input:
1625    case Fail{error}:
1626      (Tape{entries}, Fail{UnknownNode{}})
1627    case Done{input}:
1628      relu_shaped(a, entries, Core.shape(input))
1629
1630# Records relu(a), PyTorch's relu, of a's shape: the pointwise product of a
1631# copy of a's value with ReluLeft against a rank-0 +0.0, on the parallel
1632# elementwise build. The copy is saved for the VJP.
1633def relu(tape: Tape, +a: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1634  Tape{entries} = tape
1635  relu_found(a, value_counted(a, counted(entries)))
1636
1637# Records mse_loss(a, b), PyTorch's mse_loss(input, target) with mean
1638# reduction, a rank-0 tensor, from copies of the operands' values. The
1639# differences a - b and b - a, over the shapes' broadcast, are saved for the
1640# VJP. Shapes that do not broadcast are refused with IncompatibleBroadcast
1641# (6); an unknown id with UnknownNode; the tape comes back unchanged.
1642def mse_loss(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1643  binary(tape, SquaredError{}, a, b)
1644
1645def entropy_shaped(+a: Nat, +targets: List<&2, U32>, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) ->
1646  Tape & Result<&2, &2, Error, Nat>:
1647  (input, +shape) = pair
1648  loss_kept(CrossEntropy{a, shape}, entries, Loss.cross_entropy(input, targets))
1649
1650def entropy_found(+a: Nat, +targets: List<&2, U32>, pair: Found()) -> Tape & Result<&2, &2, Error, Nat>:
1651  (entries, input) = pair
1652  match input:
1653    case Fail{error}:
1654      (Tape{entries}, Fail{UnknownNode{}})
1655    case Done{input}:
1656      entropy_shaped(a, targets, entries, Core.shape(input))
1657
1658# Records cross_entropy(a, targets), PyTorch's cross_entropy with class
1659# indices and mean reduction, a rank-0 tensor, from a copy of a's value: a
1660# must be [rows, columns] logits and `targets` one class below `columns`
1661# per row. The log_softmax lp and the one-hot rows of the targets are saved
1662# for the VJP. Logits of another rank are refused with NotLogits (107),
1663# targets that do not fit with TargetOutOfRange (108), an unknown id with
1664# UnknownNode; the tape comes back unchanged.
1665def cross_entropy(tape: Tape, +a: Nat, +targets: List<&2, U32>) -> Tape & Result<&2, &2, Error, Nat>:
1666  Tape{entries} = tape
1667  entropy_found(a, targets, value_counted(a, counted(entries)))
1668
1669def error_code(error: Error) -> U32:
1670  match error:
1671    case UnknownNode{}:
1672      103
1673    case DanglingOperand{}:
1674      104
1675    case TensorError{error}:
1676      Core.error_code(error)
1677    case MatrixError{Matrix.TensorError{error}}:
1678      Core.error_code(error)
1679    case MatrixError{Matrix.ShapeError{error}}:
1680      Shape.error_code(error)
1681    case MatrixError{Matrix.RankMismatch{}}:
1682      105
1683    case MatrixError{Matrix.InnerDimensionMismatch{}}:
1684      106
1685    case TransformError{error}:
1686      Transforms.error_code(error)
1687    case BroadcastError{error}:
1688      Broadcast.error_code(error)
1689    case ReduceError{error}:
1690      Reduce.error_code(error)
1691    case GradientError{error}:
1692      Gradient.error_code(error)
1693    case LossError{error}:
1694      Loss.error_code(error)
1695    case SliceError{error}:
1696      Slice.error_code(error)
1697    case GatherError{error}:
1698      Gather.error_code(error)