LAWS.bend

bendtorch@1.3.0.0 · 325 lines · 28 definitions

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
6import Base
7import bendtorch@1.3.0.0/main.bend as T
8import bendtorch@1.3.0.0/spec.bend as Spec
9import bendtorch@1.3.0.0/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.
13law 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>.
20law 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.
28law 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.
39law 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.
51law 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 (both_sum): T.Both's own code, running two linear layers side by side on the two inputs of a pair, gives a backward
62# pass that is the adjoint of its forward pass: dy1·y1 + dy2·y2 == dx1·x1 + dx2·x2. Each layer gets its own input's
63# gradient and nothing crosses between the two sides. Attention applies RoPE to q and k this way (Both(RoPE, Both(RoPE,
64# Id))).
65law both_sum:
66  for +w1: List<&2, List<&2, Nat>>
67  for +w2: List<&2, List<&2, Nat>>
68  for +x1: List<&2, Nat>
69  for +x2: List<&2, Nat>
70  for +dy1: List<&2, Nat>
71  for +dy2: List<&2, Nat>
72  for g1: List<&2, List<&2, Nat>>
73  for g2: List<&2, List<&2, Nat>>
74  {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>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.BothL(), T.Two{w1, w2}, T.Two{x1, x2}))) == Spec.pdot(Spec.bdx(T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Layer.bwd(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.Two<List<&2, Nat>, List<&2, Nat>>, Spec.BothL(), 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>>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>, Spec.BothL(), T.Two{w1, w2}, T.Two{x1, x2})), T.Two{dy1, dy2}, T.Two{g1, g2})), T.Two{x1, x2}) : Nat}
75
76# LAW (residual): T.Residual's own code (T.ResidualOf: Seq(Fork(Id, f), Add2)), with a linear layer as f, gives a
77# backward pass that is the adjoint of its forward pass: for y = x + W·x and its dx, dy·y == dx·x. The backward pass
78# hands dy to both branches and adds what comes back, dx = dy + Wᵀ·dy: the gradient goes through the skip connection
79# untouched and through the layer, the rule every block of the Llama models relies on twice. (Over matrices the join
80# is Add2; here it is its twin over vectors, a + b forward and dy to both backward.)
81law residual:
82  for +w: List<&2, List<&2, Nat>>
83  for +x: List<&2, Nat>
84  for +dy: List<&2, Nat>
85  for g: List<&2, List<&2, Nat>>
86  {Spec.dotn(dy, Spec.fy(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, T.Layer.fwd(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, Spec.ResL(), T.Two{T.Two{Unit{}, w}, Unit{}}, x))) == Spec.dotn(Spec.bdx(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, T.Layer.bwd(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, Spec.ResL(), T.Two{T.Two{Unit{}, w}, Unit{}}, Spec.fsv(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, T.Layer.fwd(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, Spec.ResL(), T.Two{T.Two{Unit{}, w}, Unit{}}, x)), dy, T.Two{T.Two{Unit{}, g}, Unit{}})), x) : Nat}
87
88# LAW (causal_keys): T.Attention reads keys only through T.qrow, and the scores it builds for the query at position i
89# 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
90# causal) gets the i + 1 scores its condition asks for: past them, the row of probabilities is zeros.
91law causal_keys:
92  for +i: Nat
93  for +hd: Nat
94  for +qs: U32
95  for +qc: U32
96  for +ks: U32
97  for +kc: U32
98  for q: Array<F32>
99  for k: Array<F32>
100  {List.length(&2, F32, T.dr_row(T.qrow(i, hd, qs, qc, ks, kc, q, k))) == 1n+i : Nat}
101
102# LAW (causal): in T.Attention, row i of a head's probabilities is a softmax over exactly i + 1 positions. T.prow
103# writes the query's scores from the start of its row and runs the softmax over as many positions as it wrote; the count
104# 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
105# causal_keys. P starts as zeros and nothing else writes the row, so its other t - (i + 1) positions stay 0: no
106# probability lands on a later position.
107law causal:
108  for +i: Nat
109  for +t: Nat
110  for +hd: Nat
111  for +sc: F32
112  for +qs: U32
113  for +qc: U32
114  for +ks: U32
115  for +kc: U32
116  for q: Array<F32>
117  for k: Array<F32>
118  for p: Array<F32>
119  {T.pr_n(T.prow(i, t, hd, sc, qs, qc, ks, kc, q, k, p)) == 1n+i : Nat}
120
121# ---- the same three facts with every key visible: a row reads all t keys, its softmax runs over all t, and the
122# backward's G row sums over all t. The full functions take m = t - 1 (the last key's index) so that a row's count is
123# 1 + m and no subtraction appears in a statement. ----
124
125# LAW (full_keys): FullAttention's frow reads every key 0..m, so a row's scores are 1 + m long, not 1 + i.
126law full_keys:
127  for +i: Nat
128  for +m: Nat
129  for +hd: Nat
130  for +qs: U32
131  for +qc: U32
132  for +ks: U32
133  for +kc: U32
134  for q: Array<F32>
135  for k: Array<F32>
136  {List.length(&2, F32, T.dr_row(T.frow(i, m, hd, qs, qc, ks, kc, q, k))) == 1n+m : Nat}
137
138# LAW (full): row i of a head's probabilities in FullAttention is a softmax over all t positions.
139law full:
140  for +i: Nat
141  for +m: Nat
142  for +hd: Nat
143  for +sc: F32
144  for +qs: U32
145  for +qc: U32
146  for +ks: U32
147  for +kc: U32
148  for q: Array<F32>
149  for k: Array<F32>
150  for pa: Array<F32>
151  {T.pr_n(T.fprow(i, m, hd, sc, qs, qc, ks, kc, q, k, pa)) == 1n+m : Nat}
152
153# LAW (full_grad): the backward's G row for row i sums over all t keys, not 1 + i.
154law full_grad:
155  for +i: Nat
156  for +m: Nat
157  for +hd: Nat
158  for +sc: F32
159  for +ds: U32
160  for +dc: U32
161  for +vs: U32
162  for +vc: U32
163  for d: Array<F32>
164  for v: Array<F32>
165  for pa: Array<F32>
166  for ga: Array<F32>
167  {T.gr_n(T.fgrow(i, m, hd, sc, ds, dc, vs, vc, d, v, pa, ga)) == 1n+m : Nat}
168
169# LAW (causal_grad): in T.Attention's backward pass, row i of a head's G (the gradient of its scores) is computed over
170# 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
171# start of its row and runs the softmax's backward pass over as many positions as it wrote; this law says that count is
172# i + 1. G starts as zeros and nothing else writes the row, so its other t - (i + 1) positions stay 0: no gradient
173# reaches a key or a value after the query's position.
174law causal_grad:
175  for +i: Nat
176  for +t: Nat
177  for +hd: Nat
178  for +sc: F32
179  for +ds: U32
180  for +dc: U32
181  for +vs: U32
182  for +vc: U32
183  for d: Array<F32>
184  for v: Array<F32>
185  for p: Array<F32>
186  for g: Array<F32>
187  {T.gr_n(T.grow(i, t, hd, sc, ds, dc, vs, vc, d, v, p, g)) == 1n+i : Nat}
188
189# LAW (attn_par_fwd): T.Attention runs its sequences in a tree of parallel lets (T.fpar: the sequences cut in halves,
190# each half cut again, d levels deep, the leaves' outputs appended). Whatever the depth and however the sequences are
191# split between the two halves, the tree gives exactly what T.fleaf gives running them one after another, in order: the
192# parallel forward pass is the sequential one, and its output does not depend on d or on the thread count.
193law attn_par_fwd:
194  for d: Nat
195  for gs: T.Halves<T.FS>
196  for +n: Nat
197  for +h: Nat
198  for +dm: T.Dm
199  {T.fpar(d, gs, n, h, dm) == T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}) : List<&1, Array<F32>>}
200
201# LAW (attn_par_fwd_full): the same tree law for FullAttention's forward leaf: fpar_full at any depth gives exactly
202# what fleaf_full gives running the sequences one after another, so a full-attention forward does not depend on the
203# depth or the thread count either.
204law attn_par_fwd_full:
205  for d: Nat
206  for gs: T.Halves<T.FS>
207  for +n: Nat
208  for +h: Nat
209  for +dm: T.Dm
210  {T.fpar_full(d, gs, n, h, dm) == T.fleaf_full(T.hcat(T.FS, gs), h, dm, Nil{}) : List<&1, Array<F32>>}
211
212# LAW (attn_par_bwd): the same for the backward pass: T.bpar's tree gives each sequence's dQ, dK and dV exactly as
213# T.bleaf, one sequence after another.
214law attn_par_bwd:
215  for d: Nat
216  for gs: T.Halves<T.BG>
217  for +n: Nat
218  for +kv: Nat
219  for +dm: T.Dm
220  {T.bpar(d, gs, n, kv, dm) == T.bleaf(T.hcat(T.BG, gs), kv, dm, Nil{}) : List<&1, T.DB>}
221
222# LAW (attn_par_bwd_full): the tree law for FullAttention's backward leaf, as attn_par_bwd is for Attention's.
223law attn_par_bwd_full:
224  for d: Nat
225  for gs: T.Halves<T.BG>
226  for +n: Nat
227  for +kv: Nat
228  for +dm: T.Dm
229  {T.bpar_full(d, gs, n, kv, dm) == T.bleaf_full(T.hcat(T.BG, gs), kv, dm, Nil{}) : List<&1, T.DB>}
230
231# 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,
232# 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.
233law kv_of_group:
234  for +j: Nat
235  for +i: Nat
236  for +k: Nat
237  {T.kv_of(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
238
239# LAW (gqa_head): with g = 1 + i + k query heads per key/value head, query head j·g + i reads key/value head
240# (j·g + i) / g = j, the head of its own group: T.Attention's forward pass picks head hi / g, and its backward pass
241# sums the gradients of heads j·g + 0 .. j·g + g - 1 into key/value head j, so both see the same groups.
242law gqa_head:
243  for +j: Nat
244  for +i: Nat
245  for +k: Nat
246  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
247
248# LAW (kv_cols): the attention code takes the columns each query head reads from T.head_cols, in both passes (the head
249# loops pass no head index further, so they cannot compute other columns). With g = 1 + i + k query heads per key/value
250# head, query head j·g + i reads the key and value columns of head j: j·hd onwards.
251law kv_cols:
252  for +t: Nat
253  for +hd: Nat
254  for +w: Nat
255  for +kvw: Nat
256  for +sc: F32
257  for +j: Nat
258  for +i: Nat
259  for +k: Nat
260  {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}
261
262# LAW (bc_sym): broadcasting does not depend on the order of the operands: a and b broadcast to o exactly when b and a
263# do, so x + y and y + x (and x * y, y * x) accept the same shapes.
264law bc_sym:
265  for +a: List<&2, Nat>
266  for +b: List<&2, Nat>
267  for +o: List<&2, Nat>
268  {TT.bc_ok(a, b, o) == TT.bc_ok(b, a, o) : Bool}
269
270# LAW (sum_numel): summing over axis k partitions the elements: a tensor of shape s has exactly size(k) times as many
271# numbers as its sum over axis k, for every shape and every k (past the last axis the size counts as 1).
272law sum_numel:
273  for +s: List<&2, Nat>
274  for +k: Nat
275  {TT.numel(s) == Nat.mul(TT.at1(s, k), TT.numel(TT.remove(s, k))) : Nat}
276
277# LAW (narrow_rest): narrow (and setting any axis' size) changes only its own axis: with axis d taken out, the shape is
278# what it was.
279law narrow_rest:
280  for +s: List<&2, Nat>
281  for +d: Nat
282  for +l: Nat
283  {TT.remove(TT.setd(s, d, l), d) == TT.remove(s, d) : List<&2, Nat>}
284
285# LAW (narrow_axis): narrowing axis d of a tensor that has it gives that axis exactly the length asked for.
286law narrow_axis:
287  for +s: List<&2, Nat>
288  for +d: Nat
289  for +l: Nat
290  for e: {True{} == Nat.is_lt(d, TT.rank(s)) : Bool}
291  {TT.at1(TT.setd(s, d, l), d) == l : Nat}
292
293# LAW (cat_numel): concatenating along an axis the tensor has adds the counts: a tensor whose axis d is x + y long holds
294# as many numbers as the two with x and y.
295law cat_numel:
296  for +s: List<&2, Nat>
297  for +d: Nat
298  for +x: Nat
299  for +y: Nat
300  for e: {True{} == Nat.is_lt(d, TT.rank(s)) : Bool}
301  {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}
302
303# 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.
304law mat_numel:
305  for +r: Nat
306  for +c: Nat
307  {TT.numel([r, c]) == Nat.mul(r, c) : Nat}
308
309# LAW (permute_rank): permuting by p gives as many axes as p has, so a permutation of the axes keeps the rank.
310law permute_rank:
311  for +s: List<&2, Nat>
312  for p: List<&2, Nat>
313  {TT.rank(TT.permuted(s, p)) == List.length(&2, Nat, p) : Nat}
314
315# 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,
316# 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
317# 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
318# cells), both build the same tree of additions (torch/spec.bend), for every p, k and number of passes r after the
319# first p: the same pairs added in the same order, so the same bits in F32. Any block size that divides the row gives
320# the row's sum, which is why the mutant that shrinks the cross-entropy's blocks to one cell survives.
321law block_sum:
322  for +p: Nat
323  for +k: Nat
324  for +r: Nat
325  {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>}