1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/LAWS.bend as StorageLaws
4	import ../tensor/core.bend as Core
5	import ../tensor/elementwise.bend as Elementwise
6	import ../tensor/lifecycle.bend as Lifecycle
7	import ../tensor/contracts/replayed.bend as Replayed
8	import ../tensor/contracts/broadcast.bend as PointwiseBroadcast
9	import ../tensor/signed.bend as Signed
10	import ./sum_dims.bend as SumDims
11	import ./gradient.bend as Gradient
12	import ./tensor_backward.bend as Backward
13	
14	# The nodes of a list of entries, newest first.
15	def nodes_of(entries: List<&1, Backward.Entry>) -> List<&2, Backward.Node>:
16	  match entries:
17	    case Nil{}:
18	      Nil{}
19	    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
20	      node <> nodes_of(rest)
21	
22	def slots_of(entries: List<&1, Backward.Entry>) -> List<&1, Gradient.Slot>:
23	  match entries:
24	    case Nil{}:
25	      Nil{}
26	    case Con{Backward.Entry{node, shape, value, saved, slot}, rest}:
27	      slot <> slots_of(rest)
28	
29	def tape_nodes(tape: Backward.Tape) -> List<&2, Backward.Node>:
30	  Backward.Tape{entries} = tape
31	  nodes_of(entries)
32	
33	def tape_slots(tape: Backward.Tape) -> List<&1, Gradient.Slot>:
34	  Backward.Tape{entries} = tape
35	  slots_of(entries)
36	
37	# Operand bounds, stated without the engine's validator: a node may name only
38	# nodes older than itself, so each operand id must be below the number of
39	# older nodes.
40	def targets(node: Backward.Node) -> List<&2, Nat>:
41	  match node:
42	    case Backward.Leaf{}:
43	      []
44	    case Backward.Add{a, b, left, right}:
45	      [a, b]
46	    case Backward.Sub{a, b, left, right}:
47	      [a, b]
48	    case Backward.Mul{a, b, left, right}:
49	      [a, b]
50	    case Backward.Sum{a, input}:
51	      [a]
52	    case Backward.SumDims{a, input, kept}:
53	      [a]
54	    case Backward.Matmul{a, b, left, right}:
55	      [a, b]
56	    case Backward.MatmulNT{a, b, left, right}:
57	      [a, b]
58	    case Backward.Relu{a, input}:
59	      [a]
60	    case Backward.MseLoss{a, b, left, right}:
61	      [a, b]
62	    case Backward.CrossEntropy{a, input}:
63	      [a]
64	    case Backward.BatchedMatmul{a, b, left, right}:
65	      [a, b]
66	    case Backward.Slice{a, input, kept, plan}:
67	      [a]
68	    case Backward.Gathered{a, input, plan, entries}:
69	      [a]
70	    case Backward.Constant{}:
71	      []
72	    case Backward.Frozen{id, inner}:
73	      targets(inner)
74	
75	def bounded(ids: List<&2, Nat>, +bound: Nat) -> Bool:
76	  match ids:
77	    case Nil{}:
78	      True{}
79	    case Con{+id, rest}:
80	      Bool.and(Nat.is_lt(id, bound), bounded(rest, bound))
81	
82	def well_formed(nodes: List<&2, Backward.Node>) -> Bool:
83	  match nodes:
84	    case Nil{}:
85	      True{}
86	    case Con{node, +rest}:
87	      Bool.and(bounded(targets(node), List.length(&2, Backward.Node, rest)), well_formed(rest))
88	
89	# The slots newer than a target, put back in front of the result.
90	def relisted(newer: List<&1, Gradient.Slot>, result: Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>) ->
91	  Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>:
92	  match newer:
93	    case Nil{}:
94	      result
95	    case Con{slot, rest}:
96	      Backward.prepend(Gradient.Slot, slot, relisted(rest, result))
97	
98	# The entries newer than the output, each with no gradient, in front of the
99	# sweep's result.
100	def newer_unset(newer: List<&1, Backward.Entry>, swept: Backward.Swept()) -> Backward.Swept():
101	  match newer:
102	    case Nil{}:
103	      swept
104	    case Con{entry, rest}:
105	      Backward.finished(entry, Gradient.Unset{}, newer_unset(rest, swept))
106	
107	# The only contribution of a VJP, when it sends exactly one.
108	def sole(pass: Backward.Pass) -> Maybe<&1, Backward.Contribution>:
109	  Backward.Pass{saved, g, contributions} = pass
110	  match contributions:
111	    case Fail{error}:
112	      None{}
113	    case Done{Nil{}}:
114	      None{}
115	    case Done{Con{contribution, Nil{}}}:
116	      Some{contribution}
117	    case Done{Con{contribution, Con{other, rest}}}:
118	      None{}
119	
120	def sole_target(found: Maybe<&1, Backward.Contribution>) -> Maybe<&2, Nat>:
121	  match found:
122	    case None{}:
123	      None{}
124	    case Some{Backward.Contribution{target, shape, gradient}}:
125	      Some{target}
126	
127	def sole_dimensions(found: Maybe<&1, Backward.Contribution>) -> Maybe<&2, List<&2, U32>>:
128	  match found:
129	    case None{}:
130	      None{}
131	    case Some{Backward.Contribution{target, shape, Core.Packed{dimensions, tensor}}}:
132	      Some{Shape.dimensions(dimensions)}
133	
134	def sole_point(found: Maybe<&1, Backward.Contribution>, +index: U32) -> Maybe<&2, F32>:
135	  match found:
136	    case None{}:
137	      None{}
138	    case Some{Backward.Contribution{target, shape, gradient}}:
139	      Core.dynamic_read_value(Core.read_dynamic(gradient, index))
140	
141	# What a VJP that sends one contribution promises: it goes to node a, it has
142	# the target's dimensions, and every in-bounds index i holds cell(i).
143	type SoleContract<-a: Nat, -target: Shape.Shape, -cell: Nat -> F32, -pass: Backward.Pass> is Type:
144	  SoleContract{
145	    target_id: {sole_target(sole(pass)) == Some{a} : Maybe<&2, Nat>},
146	    dimensions: {sole_dimensions(sole(pass)) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>},
147	    value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
148	      {sole_point(sole(pass), index) == Some{cell(U32.to_nat(index))} : Maybe<&2, F32>}}
149	
150	# NegateLeft flips the sign bit of the left word, as PyTorch's neg does, and
151	# ignores the right word. The Sub VJP sends -g to its second operand with it.
152	law negate_left_scalar:
153	  for +left: F32
154	  for +right: F32
155	  {Elementwise.scalar(Elementwise.NegateLeft{}, left, right) == F32.neg(left) : F32}
156	
157	# A leaf sends nothing, and hands back its saved values and its gradient.
158	law tensor_vjp_leaf:
159	  for saved: Backward.Saved
160	  for g: Core.Dynamic
161	  {Backward.vjp(Backward.Leaf{}, saved, g) == Backward.Pass{saved, g, Done{Nil{}}} : Backward.Pass}
162	
163	# A constant sends nothing, and hands back its saved values and its gradient.
164	law tensor_vjp_constant:
165	  for saved: Backward.Saved
166	  for g: Core.Dynamic
167	  {Backward.vjp(Backward.Constant{}, saved, g) == Backward.Pass{saved, g, Done{Nil{}}} : Backward.Pass}
168	
169	# Frozen{id, node}, the node of a recording whose operand id is a constant,
170	# sends node's contributions without the ones to id, in order, and hands back
171	# what node's VJP hands back. Where Backward.skips(id, node) holds (node is
172	# x @ W^T and id is x and not W), the VJP does not compute the product it
173	# would drop; tensor_matmul_nt_frozen_vjp_values states what it sends.
174	law tensor_vjp_frozen_unsent:
175	  for +id: Nat
176	  for +node: Backward.Node
177	  for saved: Backward.Saved
178	  for g: Core.Dynamic
179	  for computed: {Backward.skips(id, node) == False{} : Bool}
180	  {Backward.vjp(Backward.Frozen{id, node}, saved, g) == Backward.unsent(id, Backward.vjp(node, saved, g)) : Backward.Pass}
181	
182	# Whether a list of contributions holds one to id.
183	def reaches(+id: Nat, items: List<&1, Backward.Contribution>) -> Bool:
184	  match items:
185	    case Nil{}:
186	      False{}
187	    case Con{Backward.Contribution{target, shape, gradient}, rest}:
188	      Bool.or(Nat.is_eq(id, target), reaches(id, rest))
189	
190	# Whether a VJP's pass sends a contribution to id.
191	def sends_to(+id: Nat, pass: Backward.Pass) -> Bool:
192	  Backward.Pass{saved, g, contributions} = pass
193	  match contributions:
194	    case Fail{error}:
195	      False{}
196	    case Done{items}:
197	      reaches(id, items)
198	
199	# A frozen operand gets nothing: for any node, saved values and g, no
200	# contribution of the VJP of Frozen{id, node} goes to id, skipped product
201	# included.
202	law tensor_vjp_frozen_spares:
203	  for +id: Nat
204	  for +node: Backward.Node
205	  for saved: Backward.Saved
206	  for g: Core.Dynamic
207	  {sends_to(id, Backward.vjp(Backward.Frozen{id, node}, saved, g)) == False{} : Bool}
208	
209	# backward of an output the tape does not hold is refused with UnknownNode,
210	# and the tape comes back unchanged.
211	law tensor_backward_unknown:
212	  for entries: List<&1, Backward.Entry>
213	  for +output: Nat
214	  for outside: {Nat.is_lt(output, List.length(&2, Backward.Node, nodes_of(entries))) == False{} : Bool}
215	  {Backward.backward(Backward.Tape{entries}, output) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) :
216	    Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
217	
218	# A tape with a node that refers to itself, to a newer node or to a missing
219	# node is refused with DanglingOperand, and comes back unchanged.
220	law tensor_backward_dangling:
221	  for entries: List<&1, Backward.Entry>
222	  for +output: Nat
223	  for inside: {Nat.is_lt(output, List.length(&2, Backward.Node, nodes_of(entries))) == True{} : Bool}
224	  for unsound: {Backward.wired(nodes_of(entries), List.length(&2, Backward.Node, nodes_of(entries))) == False{} : Bool}
225	  {Backward.backward(Backward.Tape{entries}, output) == (Backward.Tape{entries}, Fail{Backward.DanglingOperand{}}) :
226	    Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
227	
228	# Otherwise backward runs pass from the selected output and accumulates its
229	# working gradients into the slots.
230	law tensor_backward_success:
231	  for entries: List<&1, Backward.Entry>
232	  for +output: Nat
233	  for inside: {Nat.is_lt(output, List.length(&2, Backward.Node, nodes_of(entries))) == True{} : Bool}
234	  for sound: {Backward.wired(nodes_of(entries), List.length(&2, Backward.Node, nodes_of(entries))) == True{} : Bool}
235	  {Backward.backward(Backward.Tape{entries}, output) ==
236	    Backward.ran(Backward.pass(entries, List.length(&2, Backward.Node, nodes_of(entries)), output,
237	      Backward.at_head(List.length(&2, Backward.Node, nodes_of(entries)), output))) :
238	    Backward.Tape & Result<&2, &2, Backward.Error, Unit>}
239	
240	# The engine's wiring check is exactly the operand-bound specification above.
241	law tensor_wired_matches_operand_bounds:
242	  for +graph: List<&2, Backward.Node>
243	  {Backward.wired(graph, List.length(&2, Backward.Node, graph)) == well_formed(graph) : Bool}
244	
245	# The seed is the selected output (C1): for any entries newer than it, those
246	# entries get no gradient, and the sweep starts at the output with ones of the
247	# output's shape, Core.new(shape, 1.0), and no gradient for the older entries.
248	law tensor_seed_selected_output:
249	  for newer: List<&1, Backward.Entry>
250	  for +node: Backward.Node
251	  for +shape: Shape.Shape
252	  for value: Core.Dynamic
253	  for saved: Backward.Saved
254	  for slot: Gradient.Slot
255	  for rest: List<&1, Backward.Entry>
256	  for +older: Nat
257	  for fits: {List.length(&1, Backward.Entry, rest) == older : Nat}
258	  {Backward.pass(List.append(&1, Backward.Entry, newer, Backward.Entry{node, shape, value, saved, slot} <> rest),
259	      1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older,
260	      Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older)) ==
261	    newer_unset(newer, Backward.seeded(node, shape, value, saved, slot, rest, older, Core.new(shape, 1.0))) : Backward.Swept()}
262	
263	# The sweep starts from the seed: the output gets ones of its shape and every
264	# older entry starts with no gradient.
265	law tensor_seeded_sweep:
266	  for +node: Backward.Node
267	  for +shape: Shape.Shape
268	  for value: Core.Dynamic
269	  for saved: Backward.Saved
270	  for slot: Gradient.Slot
271	  for rest: List<&1, Backward.Entry>
272	  for +below: Nat
273	  for ones: Core.Tensor<shape>
274	  {Backward.seeded(node, shape, value, saved, slot, rest, below, Done{ones}) ==
275	    Backward.sweep(Backward.Entry{node, shape, value, saved, slot} <> rest,
276	      Done{Gradient.Held{Core.Packed{shape, ones}} <> Backward.unsets(below)}, 1n+below) : Backward.Swept()}
277	
278	# An entry that received no gradient sends nothing: the older entries are
279	# swept with their working gradients unchanged.
280	law tensor_unreached_sends_nothing:
281	  for +node: Backward.Node
282	  for +shape: Shape.Shape
283	  for value: Core.Dynamic
284	  for saved: Backward.Saved
285	  for slot: Gradient.Slot
286	  for rest: List<&1, Backward.Entry>
287	  for pending: List<&1, Gradient.Slot>
288	  for +below: Nat
289	  {Backward.sweep(Backward.Entry{node, shape, value, saved, slot} <> rest, Done{Gradient.Unset{} <> pending}, 1n+below) ==
290	    Backward.finished(Backward.Entry{node, shape, value, saved, slot}, Gradient.Unset{},
291	      Backward.sweep(rest, Done{pending}, below)) : Backward.Swept()}
292	
293	# The contributions of one VJP are delivered in list order, the first operand
294	# before the second.
295	law tensor_deliver_in_order:
296	  for +target: Nat
297	  for +shape: Shape.Shape
298	  for gradient: Core.Dynamic
299	  for rest: List<&1, Backward.Contribution>
300	  for +count: Nat
301	  for slots: List<&1, Gradient.Slot>
302	  {Backward.deliver(Backward.Contribution{target, shape, gradient} <> rest, count, Done{slots}) ==
303	    Backward.deliver(rest, count, Backward.bump(slots, count, target, shape, gradient, Backward.at_head(count, target))) :
304	    Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
305	
306	# A contribution to node target reaches exactly the slot of that node, whose
307	# Gradient.accumulate adds it to the held total (or moves it in); every other
308	# slot is kept, whatever the list holds.
309	law tensor_bump_target:
310	  for newer: List<&1, Gradient.Slot>
311	  for slot: Gradient.Slot
312	  for rest: List<&1, Gradient.Slot>
313	  for +shape: Shape.Shape
314	  for gradient: Core.Dynamic
315	  for +older: Nat
316	  for fits: {List.length(&1, Gradient.Slot, rest) == older : Nat}
317	  {Backward.bump(List.append(&1, Gradient.Slot, newer, slot <> rest), 1n+Nat.add(older, List.length(&1, Gradient.Slot, newer)),
318	      older, shape, gradient, Backward.at_head(1n+Nat.add(older, List.length(&1, Gradient.Slot, newer)), older)) ==
319	    relisted(newer, Backward.placed(rest, Gradient.accumulate(shape, slot, gradient))) :
320	    Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>}
321	
322	# After the sweep, an entry that received a gradient g accumulates it into its
323	# slot with Gradient.accumulate, so a repeated backward adds; an entry that
324	# received nothing keeps its slot.
325	law tensor_merge_accumulates:
326	  for +node: Backward.Node
327	  for +shape: Shape.Shape
328	  for value: Core.Dynamic
329	  for saved: Backward.Saved
330	  for slot: Gradient.Slot
331	  for rest: List<&1, Backward.Entry>
332	  for g: Core.Dynamic
333	  for more: List<&1, Gradient.Slot>
334	  {Backward.merge(Backward.Entry{node, shape, value, saved, slot} <> rest, Gradient.Held{g} <> more) ==
335	    Backward.merged_entry(node, shape, value, saved, Gradient.accumulate(shape, slot, g), Backward.merge(rest, more)) :
336	    List<&1, Backward.Entry> & Result<&2, &2, Backward.Error, Unit>}
337	
338	law tensor_merge_keeps_unreached:
339	  for +node: Backward.Node
340	  for +shape: Shape.Shape
341	  for value: Core.Dynamic
342	  for saved: Backward.Saved
343	  for slot: Gradient.Slot
344	  for rest: List<&1, Backward.Entry>
345	  for more: List<&1, Gradient.Slot>
346	  {Backward.merge(Backward.Entry{node, shape, value, saved, slot} <> rest, Gradient.Unset{} <> more) ==
347	    Backward.merged_entry(node, shape, value, saved, Gradient.Accumulated{slot}, Backward.merge(rest, more)) :
348	    List<&1, Backward.Entry> & Result<&2, &2, Backward.Error, Unit>}
349	
350	# zero_grad empties every slot, like PyTorch's zero_grad(set_to_none=True),
351	# and keeps every node.
352	law tensor_zero_grad_empties:
353	  for entries: List<&1, Backward.Entry>
354	  {tape_slots(Backward.zero_grad(Backward.Tape{entries})) == Backward.unsets(List.length(&1, Backward.Entry, entries)) :
355	    List<&1, Gradient.Slot>}
356	
357	law tensor_zero_grad_keeps_nodes:
358	  for entries: List<&1, Backward.Entry>
359	  {tape_nodes(Backward.zero_grad(Backward.Tape{entries})) == nodes_of(entries) : List<&2, Backward.Node>}
360	
361	# d(sum a)/da through the live broadcast_to_pointwise law: for a replayed
362	# gradient g whose shape broadcasts to exactly a's shape, the Sum VJP sends one
363	# contribution, to a, with a's dimensions, and every position i holds the
364	# gradient word at left_map(plan, i), its right-aligned broadcast projection;
365	# for the rank-0 g of a sum, that is g's only word at every position.
366	law tensor_sum_vjp_pointwise:
367	  for +a: Nat
368	  for +shape: Shape.Shape
369	  for +input: Shape.Shape
370	  for +plan: Shape.BroadcastPlan
371	  for planned: {Shape.broadcast_checked(shape, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
372	  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>}
373	  for saved: Backward.Saved
374	  for +fill: F32
375	  for +ops: List<&2, StorageLaws.Op>
376	  for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
377	  SoleContract<a, input, index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
378	    Backward.vjp(Backward.Sum{a, input}, saved, Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})>
379	
380	# d(sum of a over some axes)/da: a gradient g of `shape`, the kept or the
381	# dropped shape, with as many elements as the kept shape, is read with the
382	# kept shape and broadcast as the Sum VJP broadcasts: for a replayed g whose
383	# words are a replay over the kept shape's count, and a kept shape that
384	# broadcasts to exactly a's shape, the VJP sends one contribution, to a, with
385	# a's dimensions, and every position i holds the word of g at
386	# left_map(plan, i), the kept shape's right-aligned projection of i.
387	law tensor_sum_dims_vjp_pointwise:
388	  for +a: Nat
389	  for +shape: Shape.Shape
390	  for +input: Shape.Shape
391	  for +kept: Shape.Shape
392	  for +plan: Shape.BroadcastPlan
393	  for planned: {Shape.broadcast_checked(kept, input) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
394	  for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(input) : List<&2, U32>}
395	  for counted: {Nat.is_eq(Shape.elements(shape), Shape.elements(kept)) == True{} : Bool}
396	  for saved: Backward.Saved
397	  for +fill: F32
398	  for +ops: List<&2, StorageLaws.Op>
399	  for -fits: {Lifecycle.logical_length(Replayed.storage(kept, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}
400	  SoleContract<a, input, index => Replayed.value(kept, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
401	    Backward.vjp(Backward.SumDims{a, input, kept}, saved, Core.Packed{shape, Core.TTensor{Replayed.storage(kept, fill, ops), fits}})>
402	
403	# The tape's sum over axes takes PyTorch's signed axes: they are wrapped
404	# among the operand's axes, its rank or 1 at rank 0 (SumDims.limit), before
405	# the plan checks them. So an axis outside -limit..limit-1 is refused with
406	# AxisOutOfBounds and two axes that wrap to one with DuplicateAxis
407	# (sum_dims_LAWS.bend, signed_LAWS.bend signed_axes_bounded).
408	law tensor_sum_dims_wraps:
409	  for +a: Nat
410	  for +axes: List<&2, Signed.Position>
411	  for keepdim: Bool
412	  for entries: List<&1, Backward.Entry>
413	  for input: Core.Dynamic
414	  for +shape: Shape.Shape
415	  {Backward.dims_shaped(a, axes, keepdim, entries, (input, shape)) ==
416	    Backward.dims_planned(a, entries, input, shape, keepdim,
417	      SumDims.plan(shape, Signed.wrap_all(SumDims.limit(Shape.dimensions(shape)), axes))) :
418	    Backward.Tape & Result<&2, &2, Backward.Error, Nat>}