LAWS.bend
bendtorch@1.2.0.0 · 257 lines · 0 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.2.0.0/main.bend as T
8import bendtorch@1.2.0.0/spec.bend as Spec
9import bendtorch@1.2.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# LAW (causal_grad): in T.Attention's backward pass, row i of a head's G (the gradient of its scores) is computed over
122# 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
123# start of its row and runs the softmax's backward pass over as many positions as it wrote; this law says that count is
124# i + 1. G starts as zeros and nothing else writes the row, so its other t - (i + 1) positions stay 0: no gradient
125# reaches a key or a value after the query's position.
126law causal_grad:
127 for +i: Nat
128 for +t: Nat
129 for +hd: Nat
130 for +sc: F32
131 for +ds: U32
132 for +dc: U32
133 for +vs: U32
134 for +vc: U32
135 for d: Array<F32>
136 for v: Array<F32>
137 for p: Array<F32>
138 for g: Array<F32>
139 {T.gr_n(T.grow(i, t, hd, sc, ds, dc, vs, vc, d, v, p, g)) == 1n+i : Nat}
140
141# LAW (attn_par_fwd): T.Attention runs its sequences in a tree of parallel lets (T.fpar: the sequences cut in halves,
142# each half cut again, d levels deep, the leaves' outputs appended). Whatever the depth and however the sequences are
143# split between the two halves, the tree gives exactly what T.fleaf gives running them one after another, in order: the
144# parallel forward pass is the sequential one, and its output does not depend on d or on the thread count.
145law attn_par_fwd:
146 for d: Nat
147 for gs: T.Halves<T.FS>
148 for +n: Nat
149 for +h: Nat
150 for +dm: T.Dm
151 {T.fpar(d, gs, n, h, dm) == T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}) : List<&1, Array<F32>>}
152
153# LAW (attn_par_bwd): the same for the backward pass: T.bpar's tree gives each sequence's dQ, dK and dV exactly as
154# T.bleaf, one sequence after another.
155law attn_par_bwd:
156 for d: Nat
157 for gs: T.Halves<T.BG>
158 for +n: Nat
159 for +kv: Nat
160 for +dm: T.Dm
161 {T.bpar(d, gs, n, kv, dm) == T.bleaf(T.hcat(T.BG, gs), kv, dm, Nil{}) : List<&1, T.DB>}
162
163# 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,
164# 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.
165law kv_of_group:
166 for +j: Nat
167 for +i: Nat
168 for +k: Nat
169 {T.kv_of(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
170
171# LAW (gqa_head): with g = 1 + i + k query heads per key/value head, query head j·g + i reads key/value head
172# (j·g + i) / g = j, the head of its own group: T.Attention's forward pass picks head hi / g, and its backward pass
173# sums the gradients of heads j·g + 0 .. j·g + g - 1 into key/value head j, so both see the same groups.
174law gqa_head:
175 for +j: Nat
176 for +i: Nat
177 for +k: Nat
178 {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
179
180# LAW (kv_cols): the attention code takes the columns each query head reads from T.head_cols, in both passes (the head
181# loops pass no head index further, so they cannot compute other columns). With g = 1 + i + k query heads per key/value
182# head, query head j·g + i reads the key and value columns of head j: j·hd onwards.
183law kv_cols:
184 for +t: Nat
185 for +hd: Nat
186 for +w: Nat
187 for +kvw: Nat
188 for +sc: F32
189 for +j: Nat
190 for +i: Nat
191 for +k: Nat
192 {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}
193
194# LAW (bc_sym): broadcasting does not depend on the order of the operands: a and b broadcast to o exactly when b and a
195# do, so x + y and y + x (and x * y, y * x) accept the same shapes.
196law bc_sym:
197 for +a: List<&2, Nat>
198 for +b: List<&2, Nat>
199 for +o: List<&2, Nat>
200 {TT.bc_ok(a, b, o) == TT.bc_ok(b, a, o) : Bool}
201
202# LAW (sum_numel): summing over axis k partitions the elements: a tensor of shape s has exactly size(k) times as many
203# numbers as its sum over axis k, for every shape and every k (past the last axis the size counts as 1).
204law sum_numel:
205 for +s: List<&2, Nat>
206 for +k: Nat
207 {TT.numel(s) == Nat.mul(TT.at1(s, k), TT.numel(TT.remove(s, k))) : Nat}
208
209# LAW (narrow_rest): narrow (and setting any axis' size) changes only its own axis: with axis d taken out, the shape is
210# what it was.
211law narrow_rest:
212 for +s: List<&2, Nat>
213 for +d: Nat
214 for +l: Nat
215 {TT.remove(TT.setd(s, d, l), d) == TT.remove(s, d) : List<&2, Nat>}
216
217# LAW (narrow_axis): narrowing axis d of a tensor that has it gives that axis exactly the length asked for.
218law narrow_axis:
219 for +s: List<&2, Nat>
220 for +d: Nat
221 for +l: Nat
222 for e: {True{} == Nat.is_lt(d, TT.rank(s)) : Bool}
223 {TT.at1(TT.setd(s, d, l), d) == l : Nat}
224
225# LAW (cat_numel): concatenating along an axis the tensor has adds the counts: a tensor whose axis d is x + y long holds
226# as many numbers as the two with x and y.
227law cat_numel:
228 for +s: List<&2, Nat>
229 for +d: Nat
230 for +x: Nat
231 for +y: Nat
232 for e: {True{} == Nat.is_lt(d, TT.rank(s)) : Bool}
233 {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}
234
235# 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.
236law mat_numel:
237 for +r: Nat
238 for +c: Nat
239 {TT.numel([r, c]) == Nat.mul(r, c) : Nat}
240
241# LAW (permute_rank): permuting by p gives as many axes as p has, so a permutation of the axes keeps the rank.
242law permute_rank:
243 for +s: List<&2, Nat>
244 for p: List<&2, Nat>
245 {TT.rank(TT.permuted(s, p)) == List.length(&2, Nat, p) : Nat}
246
247# 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,
248# 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
249# 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
250# cells), both build the same tree of additions (torch/spec.bend), for every p, k and number of passes r after the
251# first p: the same pairs added in the same order, so the same bits in F32. Any block size that divides the row gives
252# the row's sum, which is why the mutant that shrinks the cross-entropy's blocks to one cell survives.
253law block_sum:
254 for +p: Nat
255 for +k: Nat
256 for +r: Nat
257 {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>}