LAWS.bend
bendtorch@1.1.1.0 · 230 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.1.1.0/main.bend as T
8import bendtorch@1.1.1.0/spec.bend as Spec
9import bendtorch@1.1.1.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 (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.
64law 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.
80law 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.
99law 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.
118law 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.
128law 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.
138law 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.
147law 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.
156law 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.
169law 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).
177law 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.
184law 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.
191law 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.
200law 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.
209law 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.
215law 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.
226law 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>}