1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/matmul.bend as Matrix
5	import ../storage/arith.bend as Arith
6	import ./broadcast.bend as Broadcast
7	import ./broadcast_pointwise_LAWS.bend as BroadcastPointwise
8	import ./broadcast_pointwise_PROOF.bend as BroadcastPointwiseProof
9	import ./gradient.bend as Gradient
10	import ./tensor_backward.bend as Backward
11	import ./tensor_backward_LAWS.bend as Laws
12	import ../tensor/contracts/replayed.bend as Replayed
13	import ../tensor/contracts/broadcast.bend as PointwiseBroadcast
14	import ../tensor/transforms.bend as Transforms
15	import ../tensor/lifecycle.bend as Lifecycle
16	import ../storage/LAWS.bend as StorageLaws
17	
18	def eq_refl(+number: Nat) -> {True{} == Nat.is_eq(number, number) : Bool}:
19	  match number:
20	    case 0n:
21	      {==}
22	    case 1n+rest:
23	      eq_refl(rest)
24	
25	def add_zero(+number: Nat) -> {number == Nat.add(number, 0n) : Nat}:
26	  match number:
27	    case 0n:
28	      {==}
29	    case 1n+rest:
30	      Equal.cong(Nat, Nat, value => 1n+value, rest, Nat.add(rest, 0n), add_zero(rest))
31	
32	def add_succ(+left: Nat, -right: Nat) -> {1n+Nat.add(left, right) == Nat.add(left, 1n+right) : Nat}:
33	  match left:
34	    case 0n:
35	      {==}
36	    case 1n+rest:
37	      Equal.cong(Nat, Nat, value => 1n+value, 1n+Nat.add(rest, right), Nat.add(rest, 1n+right), add_succ(rest, right))
38	
39	# A node with older nodes below it and extra + 1 newer ones above it never has
40	# the id of the target.
41	def older_ne(+older: Nat, -extra: Nat) -> {False{} == Nat.is_eq(Nat.add(older, 1n+extra), older) : Bool}:
42	  match older:
43	    case 0n:
44	      {==}
45	    case 1n+rest:
46	      older_ne(rest, extra)
47	
48	def listed_back(entries: List<&1, Backward.Entry>) ->
49	  {(entries, Laws.nodes_of(entries)) == Backward.listed(entries) : List<&1, Backward.Entry> & List<&2, Backward.Node>}:
50	  match entries:
51	    case Nil{}:
52	      {==}
53	    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
54	      %listed_back(rest) :
55	        {(Backward.Entry{node, shape, value, saved, slot} <> rest, node <> Laws.nodes_of(rest)) ==
56	          Backward.entry_listed(Backward.Entry{node, shape, value, saved, slot}, _) : List<&1, Backward.Entry> & List<&2, Backward.Node>}
57	      {==}
58	
59	def Laws.negate_left_scalar(left, right):
60	  {==}
61	
62	def Laws.tensor_vjp_leaf(saved, g):
63	  {==}
64	
65	def Laws.tensor_vjp_constant(saved, g):
66	  {==}
67	
68	def frozen_nt_unsent(+id: Nat, +a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Backward.Saved, g: Core.Dynamic,
69	  computed: {Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))) == False{} : Bool}) ->
70	  {Backward.nt_frozen(id, a, b, left, right, saved, g, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))) ==
71	    Backward.unsent(id, Backward.vjp_matmul_nt(a, b, left, right, saved, g)) : Backward.Pass}:
72	  %Equal.sym(Bool, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))), False{}, computed) :
73	    {Backward.nt_frozen(id, a, b, left, right, saved, g, _) == Backward.unsent(id, Backward.vjp_matmul_nt(a, b, left, right, saved, g)) : Backward.Pass}
74	  {==}
75	
76	def Laws.tensor_vjp_frozen_unsent(id, node, saved, g, computed):
77	  match node:
78	    case Backward.MatmulNT{+a, +b, +left, +right}:
79	      frozen_nt_unsent(id, a, b, left, right, saved, g, computed)
80	    case Backward.Leaf{}:
81	      {==}
82	    case Backward.Add{a, b, left, right}:
83	      {==}
84	    case Backward.Sub{a, b, left, right}:
85	      {==}
86	    case Backward.Mul{a, b, left, right}:
87	      {==}
88	    case Backward.Sum{a, input}:
89	      {==}
90	    case Backward.SumDims{a, input, kept}:
91	      {==}
92	    case Backward.Matmul{a, b, left, right}:
93	      {==}
94	    case Backward.Relu{a, input}:
95	      {==}
96	    case Backward.MseLoss{a, b, left, right}:
97	      {==}
98	    case Backward.CrossEntropy{a, input}:
99	      {==}
100	    case Backward.BatchedMatmul{a, b, left, right}:
101	      {==}
102	    case Backward.Slice{a, input, kept, plan}:
103	      {==}
104	    case Backward.Gathered{a, input, plan, entries}:
105	      {==}
106	    case Backward.Constant{}:
107	      {==}
108	    case Backward.Frozen{inner_id, inner}:
109	      {==}
110	
111	# kept_others leaves no contribution to id.
112	def spared_one(+id: Nat, +target: Nat, +shape: Shape.Shape, -gradient: Core.Dynamic, -rest: List<&1, Backward.Contribution>,
113	  ih: {Laws.reaches(id, Backward.kept_others(id, rest)) == False{} : Bool},
114	  decision: Shape.BoolCase<Nat.is_eq(id, target)>) ->
115	  {Laws.reaches(id, Backward.kept_one(target, shape, gradient, Nat.is_eq(id, target), Backward.kept_others(id, rest))) == False{} : Bool}:
116	  match decision:
117	    case Shape.BTrue{equal}:
118	      %equal : {Laws.reaches(id, Backward.kept_one(target, shape, gradient, _, Backward.kept_others(id, rest))) == False{} : Bool}
119	      ih
120	    case Shape.BFalse{+equal}:
121	      %equal : {Laws.reaches(id, Backward.kept_one(target, shape, gradient, _, Backward.kept_others(id, rest))) == False{} : Bool}
122	      %equal : {Bool.or(_, Laws.reaches(id, Backward.kept_others(id, rest))) == False{} : Bool}
123	      %Equal.sym(Bool, Laws.reaches(id, Backward.kept_others(id, rest)), False{}, ih) : {Bool.or(False{}, _) == False{} : Bool}
124	      {==}
125	
126	def spared(+id: Nat, items: List<&1, Backward.Contribution>) -> {Laws.reaches(id, Backward.kept_others(id, items)) == False{} : Bool}:
127	  match items:
128	    case Nil{}:
129	      {==}
130	    case Con{Backward.Contribution{+target, +shape, gradient}, rest}:
131	      spared_one(id, target, shape, gradient, rest, spared(id, rest), Shape.bool_case(Nat.is_eq(id, target)))
132	
133	def unsent_spared_result(+id: Nat, -saved: Backward.Saved, -g: Core.Dynamic,
134	  contributions: Result<&2, &1, Backward.Error, List<&1, Backward.Contribution>>) ->
135	  {Laws.sends_to(id, Backward.unsent(id, Backward.Pass{saved, g, contributions})) == False{} : Bool}:
136	  match contributions:
137	    case Fail{error}:
138	      {==}
139	    case Done{items}:
140	      spared(id, items)
141	
142	def unsent_spared(+id: Nat, pass: Backward.Pass) -> {Laws.sends_to(id, Backward.unsent(id, pass)) == False{} : Bool}:
143	  match pass:
144	    case Backward.Pass{saved, g, contributions}:
145	      unsent_spared_result(id, saved, g, contributions)
146	
147	def other_operand(x: Bool, y: Bool, both: {Bool.and(x, Bool.not(y)) == True{} : Bool}) -> {y == False{} : Bool}:
148	  match x:
149	    case False{}:
150	      Empty.absurd({y == False{} : Bool}, Arith.clash(both))
151	    case True{}:
152	      match y:
153	        case True{}:
154	          Empty.absurd({True{} == False{} : Bool}, Arith.clash(both))
155	        case False{}:
156	          {==}
157	
158	def alone_spared(+id: Nat, +b: Nat, +right: Shape.Shape, -x: Core.Dynamic, -y: Core.Dynamic, -g: Core.Dynamic,
159	  other: {Nat.is_eq(id, b) == False{} : Bool}, result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) ->
160	  {Laws.sends_to(id, Backward.Pass{Backward.Operands{x, y}, g, Backward.alone(Backward.multiplied(b, right, result))}) == False{} : Bool}:
161	  match result:
162	    case Fail{error}:
163	      {==}
164	    case Done{triple}:
165	      match triple:
166	        case (kept_left, rest):
167	          match rest:
168	            case (kept_right, product):
169	              %Equal.sym(Bool, Nat.is_eq(id, b), False{}, other) : {Bool.or(_, False{}) == False{} : Bool}
170	              {==}
171	
172	def right_spared(+id: Nat, +b: Nat, +right: Shape.Shape, -y: Core.Dynamic, -g: Core.Dynamic,
173	  other: {Nat.is_eq(id, b) == False{} : Bool}, copy: Core.Dynamic, pair: Core.Dynamic & Core.Dynamic) ->
174	  {Laws.sends_to(id, Backward.nt_alone_second(y, g, Backward.nt_right(b, right, copy, pair))) == False{} : Bool}:
175	  match pair:
176	    case (left, flipped):
177	      alone_spared(id, b, right, left, y, g, other, Backward.weight_product(copy, flipped))
178	
179	def first_spared(+id: Nat, +b: Nat, +right: Shape.Shape, x: Core.Dynamic, -y: Core.Dynamic,
180	  other: {Nat.is_eq(id, b) == False{} : Bool}, pair: Core.Dynamic & Core.Dynamic) ->
181	  {Laws.sends_to(id, Backward.nt_alone_first(b, right, x, y, pair)) == False{} : Bool}:
182	  match pair:
183	    case (g, copy):
184	      right_spared(id, b, right, y, g, other, copy, Core.clone_dynamic(x))
185	
186	def right_alone_spared(+id: Nat, +b: Nat, +right: Shape.Shape, saved: Backward.Saved, g: Core.Dynamic,
187	  other: {Nat.is_eq(id, b) == False{} : Bool}) ->
188	  {Laws.sends_to(id, Backward.vjp_matmul_nt_right(b, right, saved, g)) == False{} : Bool}:
189	  match saved:
190	    case Backward.Unsaved{}:
191	      {==}
192	    case Backward.Operands{x, y}:
193	      first_spared(id, b, right, x, y, other, Core.clone_dynamic(g))
194	
195	def frozen_nt_spared(+id: Nat, +a: Nat, +b: Nat, +left: Shape.Shape, +right: Shape.Shape, saved: Backward.Saved, g: Core.Dynamic,
196	  decision: Shape.BoolCase<Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))>) ->
197	  {Laws.sends_to(id, Backward.nt_frozen(id, a, b, left, right, saved, g, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))))) == False{} : Bool}:
198	  match decision:
199	    case Shape.BTrue{+equal}:
200	      %equal : {Laws.sends_to(id, Backward.nt_frozen(id, a, b, left, right, saved, g, _)) == False{} : Bool}
201	      right_alone_spared(id, b, right, saved, g,
202	        other_operand(Nat.is_eq(id, a), Nat.is_eq(id, b), Equal.sym(Bool, True{}, Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b))), equal)))
203	    case Shape.BFalse{equal}:
204	      %equal : {Laws.sends_to(id, Backward.nt_frozen(id, a, b, left, right, saved, g, _)) == False{} : Bool}
205	      unsent_spared(id, Backward.vjp_matmul_nt(a, b, left, right, saved, g))
206	
207	def Laws.tensor_vjp_frozen_spares(id, node, saved, g):
208	  match node:
209	    case Backward.MatmulNT{+a, +b, +left, +right}:
210	      frozen_nt_spared(id, a, b, left, right, saved, g, Shape.bool_case(Bool.and(Nat.is_eq(id, a), Bool.not(Nat.is_eq(id, b)))))
211	    case Backward.Leaf{}:
212	      unsent_spared(id, Backward.vjp(Backward.Leaf{}, saved, g))
213	    case Backward.Add{a, b, left, right}:
214	      unsent_spared(id, Backward.vjp(Backward.Add{a, b, left, right}, saved, g))
215	    case Backward.Sub{a, b, left, right}:
216	      unsent_spared(id, Backward.vjp(Backward.Sub{a, b, left, right}, saved, g))
217	    case Backward.Mul{a, b, left, right}:
218	      unsent_spared(id, Backward.vjp(Backward.Mul{a, b, left, right}, saved, g))
219	    case Backward.Sum{a, input}:
220	      unsent_spared(id, Backward.vjp(Backward.Sum{a, input}, saved, g))
221	    case Backward.SumDims{a, input, kept}:
222	      unsent_spared(id, Backward.vjp(Backward.SumDims{a, input, kept}, saved, g))
223	    case Backward.Matmul{a, b, left, right}:
224	      unsent_spared(id, Backward.vjp(Backward.Matmul{a, b, left, right}, saved, g))
225	    case Backward.Relu{a, input}:
226	      unsent_spared(id, Backward.vjp(Backward.Relu{a, input}, saved, g))
227	    case Backward.MseLoss{a, b, left, right}:
228	      unsent_spared(id, Backward.vjp(Backward.MseLoss{a, b, left, right}, saved, g))
229	    case Backward.CrossEntropy{a, input}:
230	      unsent_spared(id, Backward.vjp(Backward.CrossEntropy{a, input}, saved, g))
231	    case Backward.BatchedMatmul{a, b, left, right}:
232	      unsent_spared(id, Backward.vjp(Backward.BatchedMatmul{a, b, left, right}, saved, g))
233	    case Backward.Slice{a, input, kept, plan}:
234	      unsent_spared(id, Backward.vjp(Backward.Slice{a, input, kept, plan}, saved, g))
235	    case Backward.Gathered{a, input, plan, entries}:
236	      unsent_spared(id, Backward.vjp(Backward.Gathered{a, input, plan, entries}, saved, g))
237	    case Backward.Constant{}:
238	      unsent_spared(id, Backward.vjp(Backward.Constant{}, saved, g))
239	    case Backward.Frozen{inner_id, inner}:
240	      unsent_spared(id, Backward.vjp(Backward.Frozen{inner_id, inner}, saved, g))
241	
242	def Laws.tensor_backward_unknown(entries, output, outside):
243	  %listed_back(entries) :
244	    {Backward.planned(output, _) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
245	      Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
246	  %Equal.sym(Bool, Nat.is_lt(output, List.length(&2, Backward.Node, Laws.nodes_of(entries))), False{}, outside) :
247	    {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, _,
248	      Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries)))) ==
249	      (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) : Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
250	  {==}
251	
252	def Laws.tensor_backward_dangling(entries, output, inside, unsound):
253	  %listed_back(entries) :
254	    {Backward.planned(output, _) == (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) :
255	      Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
256	  %Equal.sym(Bool, Nat.is_lt(output, List.length(&2, Backward.Node, Laws.nodes_of(entries))), True{}, inside) :
257	    {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, _,
258	      Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries)))) ==
259	      (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) : Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
260	  %Equal.sym(Bool, Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries))), False{}, unsound) :
261	    {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, True{}, _) ==
262	      (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) : Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
263	  {==}
264	
265	def Laws.tensor_backward_success(entries, output, inside, sound):
266	  %listed_back(entries) :
267	    {Backward.planned(output, _) ==
268	      Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output,
269	        Backward.at_head(List.length(&2, Backward.Node, Laws.nodes_of(entries)), output))) :
270	      Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
271	  %Equal.sym(Bool, Nat.is_lt(output, List.length(&2, Backward.Node, Laws.nodes_of(entries))), True{}, inside) :
272	    {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, _,
273	      Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries)))) ==
274	      Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output,
275	        Backward.at_head(List.length(&2, Backward.Node, Laws.nodes_of(entries)), output))) :
276	      Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
277	  %Equal.sym(Bool, Backward.wired(Laws.nodes_of(entries), List.length(&2, Backward.Node, Laws.nodes_of(entries))), True{}, sound) :
278	    {Backward.checked(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output, True{}, _) ==
279	      Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, Laws.nodes_of(entries)), output,
280	        Backward.at_head(List.length(&2, Backward.Node, Laws.nodes_of(entries)), output))) :
281	      Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
282	  {==}
283	
284	def and_true(+value: Bool) -> {Bool.and(value, True{}) == value : Bool}:
285	  match value:
286	    case True{}:
287	      {==}
288	    case False{}:
289	      {==}
290	
291	def pair_bounded(+a: Nat, +b: Nat, +bound: Nat) ->
292	  {Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound)) == Laws.bounded([a, b], bound) : Bool}:
293	  Equal.cong(Bool, Bool, flag => Bool.and(Nat.is_lt(a, bound), flag), Nat.is_lt(b, bound), Bool.and(Nat.is_lt(b, bound), True{}),
294	    Equal.sym(Bool, Bool.and(Nat.is_lt(b, bound), True{}), Nat.is_lt(b, bound), and_true(Nat.is_lt(b, bound))))
295	
296	def operands_bounded(+node: Backward.Node, +bound: Nat) ->
297	  {Backward.operands_below(node, bound) == Laws.bounded(Laws.targets(node), bound) : Bool}:
298	  match node:
299	    case Backward.Leaf{}:
300	      {==}
301	    case Backward.Add{a, b, left, right}:
302	      pair_bounded(a, b, bound)
303	    case Backward.Sub{a, b, left, right}:
304	      pair_bounded(a, b, bound)
305	    case Backward.Mul{a, b, left, right}:
306	      pair_bounded(a, b, bound)
307	    case Backward.Sum{a, input}:
308	      Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
309	    case Backward.SumDims{a, input, kept}:
310	      Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
311	    case Backward.Matmul{a, b, left, right}:
312	      pair_bounded(a, b, bound)
313	    case Backward.MatmulNT{a, b, left, right}:
314	      pair_bounded(a, b, bound)
315	    case Backward.Relu{a, input}:
316	      Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
317	    case Backward.MseLoss{a, b, left, right}:
318	      pair_bounded(a, b, bound)
319	    case Backward.CrossEntropy{a, input}:
320	      Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
321	    case Backward.BatchedMatmul{a, b, left, right}:
322	      pair_bounded(a, b, bound)
323	    case Backward.Slice{a, input, kept, plan}:
324	      Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
325	    case Backward.Gathered{a, input, plan, entries}:
326	      Equal.sym(Bool, Bool.and(Nat.is_lt(a, bound), True{}), Nat.is_lt(a, bound), and_true(Nat.is_lt(a, bound)))
327	    case Backward.Constant{}:
328	      {==}
329	    case Backward.Frozen{id, inner}:
330	      operands_bounded(inner, bound)
331	
332	def Laws.tensor_wired_matches_operand_bounds(graph):
333	  match graph:
334	    case Nil{}:
335	      {==}
336	    case Con{node, rest}:
337	      %Equal.sym(Bool, Backward.operands_below(node, List.length(&2, Backward.Node, rest)),
338	        Laws.bounded(Laws.targets(node), List.length(&2, Backward.Node, rest)), operands_bounded(node, List.length(&2, Backward.Node, rest))) :
339	        {Bool.and(_, Backward.wired(rest, List.length(&2, Backward.Node, rest))) ==
340	          Bool.and(Laws.bounded(Laws.targets(node), List.length(&2, Backward.Node, rest)), Laws.well_formed(rest)) : Bool}
341	      Equal.cong(Bool, Bool, flag => Bool.and(Laws.bounded(Laws.targets(node), List.length(&2, Backward.Node, rest)), flag),
342	        Backward.wired(rest, List.length(&2, Backward.Node, rest)), Laws.well_formed(rest), Laws.tensor_wired_matches_operand_bounds(rest))
343	
344	def Laws.tensor_seed_selected_output(newer, node, shape, value, saved, slot, rest, older, fits):
345	  match newer:
346	    case Nil{}:
347	      %add_zero(older) :
348	        {Backward.pass(Backward.Entry{node, shape, value, saved, slot} <> rest, 1n+_, older, Backward.at_head(1n+_, older)) ==
349	          Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)) : Backward.Swept()}
350	      %eq_refl(older) :
351	        {Backward.pass(Backward.Entry{node, shape, value, saved, slot} <> rest, 1n+older, older, _) ==
352	          Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)) : Backward.Swept()}
353	      {==}
354	    case Con{head, tail}:
355	      %older_ne(older, List.length(&1, Backward.Entry, tail)) :
356	        {Backward.pass(head <> List.append(&1, Backward.Entry, tail, Backward.Entry{node, shape, value, saved, slot} <> rest),
357	          1n+Nat.add(older, 1n+List.length(&1, Backward.Entry, tail)), older, _) ==
358	          Backward.finished(head, Gradient.Unset{},
359	            Laws.newer_unset(tail, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)))) : Backward.Swept()}
360	      %add_succ(older, List.length(&1, Backward.Entry, tail)) :
361	        {Backward.pass(head <> List.append(&1, Backward.Entry, tail, Backward.Entry{node, shape, value, saved, slot} <> rest),
362	          1n+_, older, False{}) ==
363	          Backward.finished(head, Gradient.Unset{},
364	            Laws.newer_unset(tail, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0)))) : Backward.Swept()}
365	      Equal.cong(Backward.Swept(), Backward.Swept(), swept => Backward.finished(head, Gradient.Unset{}, swept),
366	        Backward.pass(List.append(&1, Backward.Entry, tail, Backward.Entry{node, shape, value, saved, slot} <> rest),
367	          1n+Nat.add(older, List.length(&1, Backward.Entry, tail)), older,
368	          Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, tail)), older)),
369	        Laws.newer_unset(tail, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0))),
370	        Laws.tensor_seed_selected_output(tail, node, shape, value, saved, slot, rest, older, fits))
371	
372	def Laws.tensor_seeded_sweep(node, shape, value, saved, slot, rest, below, ones):
373	  {==}
374	
375	def Laws.tensor_unreached_sends_nothing(node, shape, value, saved, slot, rest, pending, below):
376	  {==}
377	
378	def Laws.tensor_deliver_in_order(target, shape, gradient, rest, count, slots):
379	  {==}
380	
381	def Laws.tensor_bump_target(newer, slot, rest, shape, gradient, older, fits):
382	  match newer:
383	    case Nil{}:
384	      %add_zero(older) :
385	        {Backward.bump(slot <> rest, 1n+_, older, shape, gradient, Backward.at_head(1n+_, older)) ==
386	          Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)) : Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
387	      %eq_refl(older) :
388	        {Backward.bump(slot <> rest, 1n+older, older, shape, gradient, _) ==
389	          Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)) : Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
390	      {==}
391	    case Con{head, tail}:
392	      %older_ne(older, List.length(&1, Gradient.Slot, tail)) :
393	        {Backward.bump(head <> List.append(&1, Gradient.Slot, tail, slot <> rest),
394	          1n+Nat.add(older, 1n+List.length(&1, Gradient.Slot, tail)), older, shape, gradient, _) ==
395	          Backward.prepend(Gradient.Slot, head,
396	            Laws.relisted(tail, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)))) :
397	          Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
398	      %add_succ(older, List.length(&1, Gradient.Slot, tail)) :
399	        {Backward.bump(head <> List.append(&1, Gradient.Slot, tail, slot <> rest), 1n+_, older, shape, gradient, False{}) ==
400	          Backward.prepend(Gradient.Slot, head,
401	            Laws.relisted(tail, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient)))) :
402	          Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
403	      Equal.cong(Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>, Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>,
404	        result => Backward.prepend(Gradient.Slot, head, result),
405	        Backward.bump(List.append(&1, Gradient.Slot, tail, slot <> rest), 1n+Nat.add(older, List.length(&1, Gradient.Slot, tail)),
406	          older, shape, gradient, Backward.at_head(1n+Nat.add(older, List.length(&1, Gradient.Slot, tail)), older)),
407	        Laws.relisted(tail, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient))),
408	        Laws.tensor_bump_target(tail, slot, rest, shape, gradient, older, fits))
409	
410	def Laws.tensor_merge_accumulates(node, shape, value, saved, slot, rest, g, more):
411	  {==}
412	
413	def Laws.tensor_merge_keeps_unreached(node, shape, value, saved, slot, rest, more):
414	  {==}
415	
416	def Laws.tensor_zero_grad_empties(entries):
417	  match entries:
418	    case Nil{}:
419	      {==}
420	    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
421	      Equal.cong(List<&1, Gradient.Slot>, List<&1, Gradient.Slot>, slots => Gradient.Unset{} <> slots,
422	        Laws.tape_slots(Backward.zero_grad(Backward.Tape{rest})), Backward.unsets(List.length(&1, Backward.Entry, rest)),
423	        Laws.tensor_zero_grad_empties(rest))
424	
425	def Laws.tensor_zero_grad_keeps_nodes(entries):
426	  match entries:
427	    case Nil{}:
428	      {==}
429	    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
430	      Equal.cong(List<&2, Backward.Node>, List<&2, Backward.Node>, nodes => node <> nodes,
431	        Laws.tape_nodes(Backward.zero_grad(Backward.Tape{rest})), Laws.nodes_of(rest),
432	        Laws.tensor_zero_grad_keeps_nodes(rest))
433	
434	def truth(value: Bool) -> Data:
435	  match value:
436	    case True{}:
437	      Unit
438	    case False{}:
439	      Empty
440	
441	def false_true(equal: {False{} == True{} : Bool}) -> Empty:
442	  %Equal.sym(Bool, False{}, True{}, equal) : truth(_)
443	  Unit{}
444	
445	def sum_result(+a: Nat, +input: Shape.Shape, -shape: Shape.Shape, -cell: Nat -> F32, -word: U32 -> F32, -owner: Core.Dynamic,
446	  result: Result<&2, &1, Broadcast.Error, Core.Dynamic>,
447	  contract: BroadcastPointwise.BroadcastToContract<input, shape, cell, word, (owner, result)>) ->
448	  Laws.SoleContract<a, input, cell, Backward.expanded(a, input, (owner, result))>:
449	  match result:
450	    case Fail{error}:
451	      BroadcastPointwise.BroadcastToContract{successful, dimensions, value, kept} = contract
452	      Empty.absurd(Laws.SoleContract<a, input, cell, Backward.expanded(a, input, (owner, Fail{error}))>, false_true(successful))
453	    case Done{Core.Packed{output, tensor}}:
454	      BroadcastPointwise.BroadcastToContract{successful, dimensions, value, kept} = contract
455	      Laws.SoleContract{{==}, dimensions, value}
456	
457	# A successful broadcast_to hands its output to the Sum VJP's one
458	# contribution unchanged.
459	def sum_contract(+a: Nat, +input: Shape.Shape, -shape: Shape.Shape, -cell: Nat -> F32, -word: U32 -> F32,
460	  pair: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>,
461	  contract: BroadcastPointwise.BroadcastToContract<input, shape, cell, word, pair>) ->
462	  Laws.SoleContract<a, input, cell, Backward.expanded(a, input, pair)>:
463	  (owner, result) = pair
464	  sum_result(a, input, shape, cell, word, owner, result, contract)
465	
466	def Laws.tensor_sum_vjp_pointwise(a, shape, input, plan, planned, same, saved, fill, ops, fits):
467	  sum_contract(a, input, shape,
468	    index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
469	    other => Replayed.value(shape, fill, ops, other),
470	    Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, input),
471	    BroadcastPointwise.broadcast_to_pointwise(shape, input, plan, planned, same, fill, ops, fits))
472	
473	# restored only puts g back in its own shape: the contributions are the pass's.
474	def restored_contract(-a: Nat, -target: Shape.Shape, -cell: Nat -> F32, +shape: Shape.Shape, pass: Backward.Pass,
475	  contract: Laws.SoleContract<a, target, cell, pass>) -> Laws.SoleContract<a, target, cell, Backward.restored(shape, pass)>:
476	  match pass:
477	    case Backward.Pass{saved, g, Fail{error}}:
478	      Laws.SoleContract{target_id, dimensions, value} = contract
479	      Laws.SoleContract{target_id, dimensions, value}
480	    case Backward.Pass{saved, g, Done{Nil{}}}:
481	      Laws.SoleContract{target_id, dimensions, value} = contract
482	      Laws.SoleContract{target_id, dimensions, value}
483	    case Backward.Pass{saved, g, Done{Con{contribution, Nil{}}}}:
484	      Laws.SoleContract{target_id, dimensions, value} = contract
485	      Laws.SoleContract{target_id, dimensions, value}
486	    case Backward.Pass{saved, g, Done{Con{contribution, Con{other, rest}}}}:
487	      Laws.SoleContract{target_id, dimensions, value} = contract
488	      Laws.SoleContract{target_id, dimensions, value}
489	
490	def sum_dims_decided(+a: Nat, +shape: Shape.Shape, +input: Shape.Shape, +kept: Shape.Shape, +plan: Shape.BroadcastPlan,
491	  planned: {Shape.broadcast_checked(kept, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
492	  same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>},
493	  counted: {Nat.is_eq(Shape.elements(shape), Shape.elements(kept)) == True{} : Bool},
494	  saved: Backward.Saved, +fill: F32, +ops: List<&2, StorageLaws.Op>,
495	  -fits: {Lifecycle.logical_length(Replayed.storage(kept, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
496	  decision: Shape.BoolCase<Nat.is_eq(Shape.elements(shape), Shape.elements(kept))>) ->
497	  Laws.SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
498	    Backward.unsqueezed(a, input, kept, shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}, decision)>:
499	  match decision:
500	    case Shape.BFalse{equal}:
501	      Empty.absurd(Laws.SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
502	        Backward.unsqueezed(a, input, kept, shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}, Shape.BFalse{equal})>,
503	        Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(Shape.elements(shape), Shape.elements(kept)), True{}, equal, counted)))
504	    case Shape.BTrue{equal}:
505	      restored_contract(a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))), shape,
506	        Backward.vjp_sum(a, input, Transforms.reshape_equal(shape, kept, Core.TTensor{Replayed.storage(kept, fill, ops), fits},
507	          Transforms.elements_equal(shape, kept, equal))),
508	        Laws.tensor_sum_vjp_pointwise(a, kept, input, plan, planned, same, saved, fill, ops,
509	          Equal.trans(Nat, Lifecycle.logical_length(Replayed.storage(kept, fill, ops)), Shape.logical_numel(Shape.dimensions(shape)),
510	            Shape.logical_numel(Shape.dimensions(kept)), fits, Transforms.elements_equal(shape, kept, equal))))
511	
512	def Laws.tensor_sum_dims_vjp_pointwise(a, shape, input, kept, plan, planned, same, counted, saved, fill, ops, fits):
513	  sum_dims_decided(a, shape, input, kept, plan, planned, same, counted, saved, fill, ops, fits,
514	    Shape.bool_case(Nat.is_eq(Shape.elements(shape), Shape.elements(kept))))
515	
516	def Laws.tensor_sum_dims_wraps(a, axes, keepdim, entries, input, shape):
517	  {==}