1	# BendTorch's laws. Each is a claim about every input, not a test case; PROOF.bend proves them, and
2	# `bend PROOF.bend --verdict` checks the proofs again with Bend's kernel proved in Lean. Plain-language explanations
3	# are in torch/README.md. The laws about linear maps speak about the exact model in torch/spec.bend; the others speak
4	# about the package's own functions.
5	
6	import Base
7	import ./main.bend as T
8	import ./spec.bend as Spec
9	import ./tensor.bend as TT
10	
11	# LAW (adjoint): the backward pass of a linear layer, dx = Wᵀ·dy, is the adjoint of its forward pass y = W·x: for any
12	# W, dy and x, dy·(W·x) == (Wᵀ·dy)·x. This is what makes dx the gradient of the loss with respect to x.
13	law adjoint:
14	  for w: List<&2, List<&2, Nat>>
15	  for g: List<&2, Nat>
16	  for +v: List<&2, Nat>
17	  {Spec.dotn(g, Spec.mv(w, v)) == Spec.dotn(Spec.mtv(w, g), v) : Nat}
18	
19	# LAW (weight_adjoint): the weight gradient dW = dy ⊗ x is the adjoint in W: dy·(W·x) == <dy ⊗ x, W>.
20	law weight_adjoint:
21	  for w: List<&2, List<&2, Nat>>
22	  for g: List<&2, Nat>
23	  for +v: List<&2, Nat>
24	  {Spec.dotn(g, Spec.mv(w, v)) == Spec.frob(Spec.outer(g, v), w) : Nat}
25	
26	# LAW (adjoint_seq): Seq runs the second layer's backward pass, then the first's. For two linear layers that is
27	# W1ᵀ·(W2ᵀ·dy), and it is the adjoint of the composed forward pass W2·(W1·x): the chain rule, for linear maps.
28	law adjoint_seq:
29	  for +w1: List<&2, List<&2, Nat>>
30	  for +w2: List<&2, List<&2, Nat>>
31	  for +g: List<&2, Nat>
32	  for +v: List<&2, Nat>
33	  {Spec.dotn(g, Spec.mv(w2, Spec.mv(w1, v))) == Spec.dotn(Spec.mtv(w1, Spec.mtv(w2, g)), v) : Nat}
34	
35	# LAW (seq_chain): T.Seq's own code, running two linear layers of the exact model one after the other, gives a backward
36	# pass that is the adjoint of its forward pass: dy·y == dx·x, for the y the forward pass returns and the dx its backward
37	# pass returns from the state the forward pass saved. adjoint_seq says it of the formulas; this says it of the
38	# combinator that composes the layers: Seq feeds the second layer's input gradient to the first, the chain rule.
39	law seq_chain:
40	  for +w1: List<&2, List<&2, Nat>>
41	  for +w2: List<&2, List<&2, Nat>>
42	  for +x: List<&2, Nat>
43	  for +dy: List<&2, Nat>
44	  for g1: List<&2, List<&2, Nat>>
45	  for g2: List<&2, List<&2, Nat>>
46	  {Spec.dotn(dy, Spec.fy(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.fwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.SeqL(), T.Two{w1, w2}, x))) == Spec.dotn(Spec.bdx(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Layer.bwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.SeqL(), T.Two{w1, w2}, Spec.fsv(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.fwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.SeqL(), T.Two{w1, w2}, x)), dy, T.Two{g1, g2})), x) : Nat}
47	
48	# LAW (fork_sum): T.Fork's own code, giving one input to two linear layers, gives as input gradient the sum of the two
49	# layers' input gradients, and that sum is the adjoint of the pair of outputs: dy1·y1 + dy2·y2 == dx·x. A residual
50	# block (Fork of the identity and a layer) and attention's q, k, v projections rely on it.
51	law fork_sum:
52	  for +w1: List<&2, List<&2, Nat>>
53	  for +w2: List<&2, List<&2, Nat>>
54	  for +x: List<&2, Nat>
55	  for +dy1: List<&2, Nat>
56	  for +dy2: List<&2, Nat>
57	  for g1: List<&2, List<&2, Nat>>
58	  for g2: List<&2, List<&2, Nat>>
59	  {Spec.pdot(T.Two{dy1, dy2}, Spec.fy(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.fwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.ForkL(), T.Two{w1, w2}, x))) == Spec.dotn(Spec.bdx(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Layer.bwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.ForkL(), T.Two{w1, w2}, Spec.fsv(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.fwd(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.ForkL(), T.Two{w1, w2}, x)), T.Two{dy1, dy2}, T.Two{g1, g2})), x) : Nat}
60	
61	# LAW (causal_keys): T.Attention reads keys only through T.qrow, and the scores it builds for the query at position i
62	# come from keys 0..i, one per key: exactly i + 1 of them. So no key after position i is read, and T.causal_row (law
63	# causal) gets the i + 1 scores its condition asks for: past them, the row of probabilities is zeros.
64	law causal_keys:
65	  for +i: Nat
66	  for +hd: Nat
67	  for +qs: U32
68	  for +qc: U32
69	  for +ks: U32
70	  for +kc: U32
71	  for q: Array<F32>
72	  for k: Array<F32>
73	  {List.length(&2, F32, T.dr_row(T.qrow(i, hd, qs, qc, ks, kc, q, k))) == 1n+i : Nat}
74	
75	# LAW (causal): in T.Attention, row i of a head's probabilities is a softmax over exactly i + 1 positions. T.prow
76	# writes the query's scores from the start of its row and runs the softmax over as many positions as it wrote; the count
77	# is the one the softmax's own division loop reports (T.psmc), and this law says it is i + 1, the keys 0..i of
78	# causal_keys. P starts as zeros and nothing else writes the row, so its other t - (i + 1) positions stay 0: no
79	# probability lands on a later position.
80	law causal:
81	  for +i: Nat
82	  for +t: Nat
83	  for +hd: Nat
84	  for +sc: F32
85	  for +qs: U32
86	  for +qc: U32
87	  for +ks: U32
88	  for +kc: U32
89	  for q: Array<F32>
90	  for k: Array<F32>
91	  for p: Array<F32>
92	  {T.pr_n(T.prow(i, t, hd, sc, qs, qc, ks, kc, q, k, p)) == 1n+i : Nat}
93	
94	# LAW (causal_grad): in T.Attention's backward pass, row i of a head's G (the gradient of its scores) is computed over
95	# exactly i + 1 positions. T.grow reads dP's row against values 0..i only (as qrow reads keys 0..i), writes it from the
96	# start of its row and runs the softmax's backward pass over as many positions as it wrote; this law says that count is
97	# i + 1. G starts as zeros and nothing else writes the row, so its other t - (i + 1) positions stay 0: no gradient
98	# reaches a key or a value after the query's position.
99	law causal_grad:
100	  for +i: Nat
101	  for +t: Nat
102	  for +hd: Nat
103	  for +sc: F32
104	  for +ds: U32
105	  for +dc: U32
106	  for +vs: U32
107	  for +vc: U32
108	  for d: Array<F32>
109	  for v: Array<F32>
110	  for p: Array<F32>
111	  for g: Array<F32>
112	  {T.gr_n(T.grow(i, t, hd, sc, ds, dc, vs, vc, d, v, p, g)) == 1n+i : Nat}
113	
114	# LAW (attn_par_fwd): T.Attention runs its sequences in a tree of parallel lets (T.fpar: the sequences cut in halves,
115	# each half cut again, d levels deep, the leaves' outputs appended). Whatever the depth and however the sequences are
116	# split between the two halves, the tree gives exactly what T.fleaf gives running them one after another, in order: the
117	# parallel forward pass is the sequential one, and its output does not depend on d or on the thread count.
118	law attn_par_fwd:
119	  for d: Nat
120	  for gs: T.Halves<T.FS>
121	  for +n: Nat
122	  for +h: Nat
123	  for +dm: T.Dm
124	  {T.fpar(d, gs, n, h, dm) == T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}) : List<&1, Array<F32>>}
125	
126	# LAW (attn_par_bwd): the same for the backward pass: T.bpar's tree gives each sequence's dQ, dK and dV exactly as
127	# T.bleaf, one sequence after another.
128	law attn_par_bwd:
129	  for d: Nat
130	  for gs: T.Halves<T.BG>
131	  for +n: Nat
132	  for +kv: Nat
133	  for +dm: T.Dm
134	  {T.bpar(d, gs, n, kv, dm) == T.bleaf(T.hcat(T.BG, gs), kv, dm, Nil{}) : List<&1, T.DB>}
135	
136	# LAW (kv_of_group): T.kv_of, the function both passes of T.Attention use to pick the key/value head of query head hi,
137	# gives j for every head of group j: query heads j·g + 0 .. j·g + g - 1 (g = 1 + i + k, i < g) all read key/value head j.
138	law kv_of_group:
139	  for +j: Nat
140	  for +i: Nat
141	  for +k: Nat
142	  {T.kv_of(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
143	
144	# LAW (gqa_head): with g = 1 + i + k query heads per key/value head, query head j·g + i reads key/value head
145	# (j·g + i) / g = j, the head of its own group: T.Attention's forward pass picks head hi / g, and its backward pass
146	# sums the gradients of heads j·g + 0 .. j·g + g - 1 into key/value head j, so both see the same groups.
147	law gqa_head:
148	  for +j: Nat
149	  for +i: Nat
150	  for +k: Nat
151	  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
152	
153	# LAW (kv_cols): the attention code takes the columns each query head reads from T.head_cols, in both passes (the head
154	# loops pass no head index further, so they cannot compute other columns). With g = 1 + i + k query heads per key/value
155	# head, query head j·g + i reads the key and value columns of head j: j·hd onwards.
156	law kv_cols:
157	  for +t: Nat
158	  for +hd: Nat
159	  for +w: Nat
160	  for +kvw: Nat
161	  for +sc: F32
162	  for +j: Nat
163	  for +i: Nat
164	  for +k: Nat
165	  {T.hc_kc(T.head_cols(T.Dm{t, 1n+Nat.add(i, k), hd, w, kvw, sc}, Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i))) == U32.from_nat(Nat.mul(j, hd)) : U32}
166	
167	# LAW (bc_sym): broadcasting does not depend on the order of the operands: a and b broadcast to o exactly when b and a
168	# do, so x + y and y + x (and x * y, y * x) accept the same shapes.
169	law bc_sym:
170	  for +a: List<&2, Nat>
171	  for +b: List<&2, Nat>
172	  for +o: List<&2, Nat>
173	  {TT.bc_ok(a, b, o) == TT.bc_ok(b, a, o) : Bool}
174	
175	# LAW (sum_numel): summing over axis k partitions the elements: a tensor of shape s has exactly size(k) times as many
176	# numbers as its sum over axis k, for every shape and every k (past the last axis the size counts as 1).
177	law sum_numel:
178	  for +s: List<&2, Nat>
179	  for +k: Nat
180	  {TT.numel(s) == Nat.mul(TT.at1(s, k), TT.numel(TT.remove(s, k))) : Nat}
181	
182	# LAW (narrow_rest): narrow (and setting any axis' size) changes only its own axis: with axis d taken out, the shape is
183	# what it was.
184	law narrow_rest:
185	  for +s: List<&2, Nat>
186	  for +d: Nat
187	  for +l: Nat
188	  {TT.remove(TT.setd(s, d, l), d) == TT.remove(s, d) : List<&2, Nat>}
189	
190	# LAW (narrow_axis): narrowing axis d of a tensor that has it gives that axis exactly the length asked for.
191	law narrow_axis:
192	  for +s: List<&2, Nat>
193	  for +d: Nat
194	  for +l: Nat
195	  for e: {True{} == Nat.is_lt(d, TT.rank(s)) : Bool}
196	  {TT.at1(TT.setd(s, d, l), d) == l : Nat}
197	
198	# LAW (cat_numel): concatenating along an axis the tensor has adds the counts: a tensor whose axis d is x + y long holds
199	# as many numbers as the two with x and y.
200	law cat_numel:
201	  for +s: List<&2, Nat>
202	  for +d: Nat
203	  for +x: Nat
204	  for +y: Nat
205	  for e: {True{} == Nat.is_lt(d, TT.rank(s)) : Bool}
206	  {TT.numel(TT.setd(s, d, Nat.add(x, y))) == Nat.add(TT.numel(TT.setd(s, d, x)), TT.numel(TT.setd(s, d, y))) : Nat}
207	
208	# LAW (mat_numel): a Tensor<[r, c]> holds the r·c numbers of a Mat<r, c>, so of_mat and to_mat reuse the array as it is.
209	law mat_numel:
210	  for +r: Nat
211	  for +c: Nat
212	  {TT.numel([r, c]) == Nat.mul(r, c) : Nat}
213	
214	# LAW (permute_rank): permuting by p gives as many axes as p has, so a permutation of the axes keeps the rank.
215	law permute_rank:
216	  for +s: List<&2, Nat>
217	  for p: List<&2, Nat>
218	  {TT.rank(TT.permuted(s, p)) == List.length(&2, Nat, p) : Nat}
219	
220	# LAW (block_sum): summing a row in blocks gives the row's own sum, bit for bit. tree_sum adds a row's cells in pairs,
221	# pass after pass. The GPU's loss kernels instead give each block of 2^p cells to one leaf, which sums it as a balanced
222	# tree of pairs, and then run tree_sum's remaining passes over the block sums. When 2^p divides the row (a row of 2^p·k
223	# cells), both build the same tree of additions (torch/spec.bend), for every p, k and number of passes r after the
224	# first p: the same pairs added in the same order, so the same bits in F32. Any block size that divides the row gives
225	# the row's sum, which is why the mutant that shrinks the cross-entropy's blocks to one cell survives.
226	law block_sum:
227	  for +p: Nat
228	  for +k: Nat
229	  for +r: Nat
230	  {Spec.passes(Nat.add(r, p), Spec.blocks(0n, Nat.mul(Spec.pw(p), k), 0n)) == Spec.passes(r, Spec.blocks(p, k, 0n)) : List<&2, Spec.Sum>}