1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/elementwise.bend as Elementwise
5	import ../tensor/layout.bend as Layout
6	import ../tensor/matmul.bend as Matrix
7	import ../tensor/batched_matmul.bend as Batched
8	import ../tensor/reduction.bend as Reduction
9	import ../tensor/transforms.bend as Transforms
10	import ../tensor/signed.bend as Signed
11	import ../tensor/slice.bend as Slice
12	import ../tensor/gather.bend as Gather
13	import ./broadcast.bend as Broadcast
14	import ./reduce.bend as Reduce
15	import ./sum_to_size.bend as SumToSize
16	import ./sum_dims.bend as SumDims
17	import ./gradient.bend as Gradient
18	import ./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	
69	type 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	
82	type 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.
101	type Saved is Type:
102	  Unsaved{}
103	  Operands{left: Core.Dynamic, right: Core.Dynamic}
104	
105	type Entry is Type:
106	  Entry{node: Node, shape: Shape.Shape, value: Core.Dynamic, saved: Saved, grad: Gradient.Slot}
107	
108	type Tape is Type:
109	  Tape{entries: List<&1, Entry>}
110	
111	type 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.
116	type Pass is Type:
117	  Pass{saved: Saved, gradient: Core.Dynamic, contributions: Result<&2, &1, Error, List<&1, Contribution>>}
118	
119	type Binary is Data:
120	  Plus{}
121	  Minus{}
122	  Times{}
123	  Product{}
124	  ProductNT{}
125	  SquaredError{}
126	  BatchedProduct{}
127	
128	def empty() -> Tape:
129	  Tape{Nil{}}
130	
131	# Whether the head of a newest-first list of count elements has id `id`.
132	def 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	
139	def 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.
145	def 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	
152	def 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	
159	def 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	
168	def 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	
178	def 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.
189	def 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	
192	def contribution_of(pair: Core.Dynamic & Result<&2, &1, Error, Contribution>) -> Result<&2, &1, Error, Contribution>:
193	  (owner, result) = pair
194	  result
195	
196	def 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	
204	def 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	
208	def 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.
213	def 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	
216	def 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.
223	def 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	
231	def negate(g: Core.Dynamic) -> Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
232	  negated(g, Core.new(Reduction.scalar_shape(), 0.0))
233	
234	def 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	
239	def 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.
244	def 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	
247	def 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	
253	def 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.
261	def 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	
268	def 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).
277	def vjp_sum(+a: Nat, +input: Shape.Shape, g: Core.Dynamic) -> Pass:
278	  expanded(a, input, Broadcast.broadcast_to(g, input))
279	
280	def 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.
290	def 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.
295	def restored(+shape: Shape.Shape, pass: Pass) -> Pass:
296	  Pass{saved, gradient, contributions} = pass
297	  Pass{saved, relabeled(shape, gradient), contributions}
298	
299	def 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.
310	def 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	
314	def 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.
325	def 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.
335	def 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	
344	def 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	
349	def 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	
354	def 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	
359	def 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.
366	def 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.
375	def 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.
383	def 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.
387	def 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	
392	def 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	
397	def 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	
402	def 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	
407	def 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	
412	def 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	
419	def 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	
423	def 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.
430	def 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.
441	def 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	
448	def 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.
457	def 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	
464	def 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	
470	def 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	
476	def 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.
486	def 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	
493	def 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	
499	def 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.
514	def 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	
522	def 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.
531	def 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	
538	def 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.
546	def 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.
552	def 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	
559	def 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.
569	def 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	
578	def 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.
585	def 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	
590	def 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	
595	def 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.
602	def 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	
607	def 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	
612	def 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.
626	def 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	
633	def 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	
641	def 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.
654	def 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	
658	def 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.
670	def 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.
677	def 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	
718	def 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.
728	def 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	
740	def 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	
750	def 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	
760	def Swept() -> Type:
761	  List<&1, Entry> & Result<&2, &1, Error, List<&1, Gradient.Slot>>
762	
763	def finished(entry: Entry, slot: Gradient.Slot, pair: Swept()) -> Swept():
764	  (rest, result) = pair
765	  (entry <> rest, prepend(Gradient.Slot, slot, result))
766	
767	def 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	
772	def 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.
782	def 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	
795	def 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.
805	def 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.
812	def 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	
823	def 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	
858	def 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.
869	def 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	
879	def 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	
890	def 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	
894	def 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	
902	def 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	
911	def 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.
920	def backward(tape: Tape, +output: Nat) -> Tape & Result<&2, &2, Error, Unit>:
921	  Tape{entries} = tape
922	  planned(output, listed(entries))
923	
924	def 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).
932	def zero_grad(tape: Tape) -> Tape:
933	  Tape{entries} = tape
934	  Tape{cleared(entries)}
935	
936	# Reads.
937	
938	def Found() -> Type:
939	  List<&1, Entry> & Result<&2, &1, Error, Core.Dynamic>
940	
941	def 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	
946	def 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	
950	def 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.
955	def 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	
966	def 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.
971	def 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	
978	def 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	
982	def 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.
987	def 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	
991	def Graded() -> Type:
992	  List<&1, Entry> & Result<&2, &1, Error, Maybe<&1, Core.Dynamic>>
993	
994	def held_copy(pair: Core.Dynamic & Core.Dynamic) -> Gradient.Slot & Maybe<&1, Core.Dynamic>:
995	  (gradient, copy) = pair
996	  (Gradient.Held{gradient}, Some{copy})
997	
998	def 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	
1005	def 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	
1010	def 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	
1014	def prepend_graded(entry: Entry, pair: Graded()) -> Graded():
1015	  (rest, result) = pair
1016	  (entry <> rest, result)
1017	
1018	def 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	
1029	def 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	
1033	def 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.
1038	def 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	
1044	def 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	
1049	def done_id(pair: Tape & Nat) -> Tape & Result<&2, &2, Error, Nat>:
1050	  (tape, id) = pair
1051	  (tape, Done{id})
1052	
1053	def 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.
1061	def 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.
1073	def 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	
1108	def 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.
1117	def 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.
1126	def 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.
1136	def 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	
1143	def 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.
1152	def 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.
1156	def 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.
1164	def constant(tape: Tape, value: Core.Dynamic) -> Tape & Nat:
1165	  Tape{entries} = tape
1166	  pushed(Constant{}, value, Unsaved{}, counted(entries))
1167	
1168	def 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	
1176	def 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	
1185	def 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	
1190	def 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	
1196	def 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	
1204	def 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	
1209	def 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`.
1216	def 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	
1224	def 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	
1243	def 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	
1249	def 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	
1257	def 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.
1269	def 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	
1273	def add(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1274	  binary(tape, Plus{}, a, b)
1275	
1276	def sub(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1277	  binary(tape, Minus{}, a, b)
1278	
1279	def mul(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1280	  binary(tape, Times{}, a, b)
1281	
1282	def 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.
1290	def 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.
1295	def 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.
1299	def 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.
1309	def 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	
1317	def 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	
1325	def 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	
1332	def 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.
1343	def 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	
1347	def 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	
1356	def 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	
1360	def 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.
1369	def 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	
1373	def 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	
1380	def 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	
1389	def 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.
1399	def 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	
1404	def 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.
1419	def 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	
1423	def 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	
1431	def 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	
1439	def 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).
1447	type 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	
1452	def slice_extent(+shape: Shape.Shape, +axis: U32) -> U32:
1453	  Signed.extent(Shape.dimensions(shape), axis)
1454	
1455	def 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	
1466	def 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.
1479	def 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	
1482	def 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	
1491	def 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	
1495	def 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.
1510	def 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`.
1515	def 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).
1520	def 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.
1524	def 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	
1527	def 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.
1537	type Indexing is Data:
1538	  IndexSelected{dim: Signed.Position, entries: List<&2, U32>}
1539	  GatheredBy{dim: Signed.Position, index: Gather.Index}
1540	
1541	def 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	
1548	def 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	
1555	def 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	
1562	def 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	
1571	def 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	
1576	def 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	
1581	def 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.
1597	def 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).
1602	def 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).
1606	def 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	
1609	def 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	
1618	def 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	
1622	def 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.
1633	def 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.
1642	def mse_loss(tape: Tape, +a: Nat, +b: Nat) -> Tape & Result<&2, &2, Error, Nat>:
1643	  binary(tape, SquaredError{}, a, b)
1644	
1645	def 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	
1650	def 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.
1665	def 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	
1669	def 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)