PROOF.bend
bendtorch@1.0.0.0 · 779 lines · 0 definitions
1# The proofs of LAWS.bend. Every proof is plain Bend: no @unsafe, no ?TODO, no axioms. The arithmetic lemmas come
2# from bend-ml-nat-lemmas; the division lemmas behind gqa_head are llama.bend's (MIT, same author).
3
4import Base
5import bendtorch@1.0.0.0/main.bend as T
6import bendtorch@1.0.0.0/spec.bend as Spec
7import bendtorch@1.0.0.0/tensor.bend as TT
8import bendtorch@1.0.0.0/LAWS.bend as Laws
9import bend-ml-nat-lemmas@0.1.1.0/main.bend as NL
10import bend-ml-tensor-array@0.1.14.0/main.bend as TA
11
12# =====================================================================================================
13# Arithmetic
14# =====================================================================================================
15
16# (a + b) + (c + d) == (a + c) + (b + d)
17def add4(+a: Nat, +b: Nat, +c: Nat, +d: Nat) -> {Nat.add(Nat.add(a, b), Nat.add(c, d)) == Nat.add(Nat.add(a, c), Nat.add(b, d)) : Nat}:
18 Equal.trans(Nat, Nat.add(Nat.add(a, b), Nat.add(c, d)), Nat.add(a, Nat.add(b, Nat.add(c, d))), Nat.add(Nat.add(a, c), Nat.add(b, d)),
19 Equal.sym(Nat, Nat.add(a, Nat.add(b, Nat.add(c, d))), Nat.add(Nat.add(a, b), Nat.add(c, d)), NL.add_assoc(a, b, Nat.add(c, d))),
20 Equal.trans(Nat, Nat.add(a, Nat.add(b, Nat.add(c, d))), Nat.add(a, Nat.add(c, Nat.add(b, d))), Nat.add(Nat.add(a, c), Nat.add(b, d)),
21 Equal.cong(Nat, Nat, s => Nat.add(a, s), Nat.add(b, Nat.add(c, d)), Nat.add(c, Nat.add(b, d)), NL.add_swap(b, c, d)),
22 NL.add_assoc(a, c, Nat.add(b, d))))
23
24# k * (a + b) == k * a + k * b
25def mul_dist_l(+k: Nat, +a: Nat, +b: Nat) -> {Nat.mul(k, Nat.add(a, b)) == Nat.add(Nat.mul(k, a), Nat.mul(k, b)) : Nat}:
26 Equal.trans(Nat, Nat.mul(k, Nat.add(a, b)), Nat.mul(Nat.add(a, b), k), Nat.add(Nat.mul(k, a), Nat.mul(k, b)),
27 NL.mul_comm(k, Nat.add(a, b)),
28 Equal.trans(Nat, Nat.mul(Nat.add(a, b), k), Nat.add(Nat.mul(a, k), Nat.mul(b, k)), Nat.add(Nat.mul(k, a), Nat.mul(k, b)),
29 Equal.sym(Nat, Nat.add(Nat.mul(a, k), Nat.mul(b, k)), Nat.mul(Nat.add(a, b), k), NL.mul_dist(a, b, k)),
30 Equal.trans(Nat, Nat.add(Nat.mul(a, k), Nat.mul(b, k)), Nat.add(Nat.mul(k, a), Nat.mul(b, k)), Nat.add(Nat.mul(k, a), Nat.mul(k, b)),
31 Equal.cong(Nat, Nat, s => Nat.add(s, Nat.mul(b, k)), Nat.mul(a, k), Nat.mul(k, a), NL.mul_comm(a, k)),
32 Equal.cong(Nat, Nat, s => Nat.add(Nat.mul(k, a), s), Nat.mul(b, k), Nat.mul(k, b), NL.mul_comm(b, k)))))
33
34# =====================================================================================================
35# Dot products
36# =====================================================================================================
37
38# dot(a + b, v) == dot(a, v) + dot(b, v): induction on the three vectors; at each entry (x + y)·z splits into x·z + y·z
39# (distributivity) and add4 regroups the four terms
40def dot_vadd(a: List<&2, Nat>, b: List<&2, Nat>, v: List<&2, Nat>) -> {Spec.dotn(Spec.vadd(a, b), v) == Nat.add(Spec.dotn(a, v), Spec.dotn(b, v)) : Nat}:
41 match a b v:
42 case Nil{} _ _:
43 {==}
44 case Con{+x, +xt} Nil{} +v:
45 Equal.sym(Nat, Nat.add(Spec.dotn(Con{x, xt}, v), 0n), Spec.dotn(Con{x, xt}, v), NL.add_zero(Spec.dotn(Con{x, xt}, v)))
46 case Con{x, xt} Con{y, yt} Nil{}:
47 {==}
48 case Con{+x, +xt} Con{+y, +yt} Con{+z, +zt}:
49 +xz = Nat.mul(x, z)
50 +yz = Nat.mul(y, z)
51 +pa = Spec.dotn(xt, zt)
52 +pb = Spec.dotn(yt, zt)
53 Equal.trans(Nat, Nat.add(Nat.mul(Nat.add(x, y), z), Spec.dotn(Spec.vadd(xt, yt), zt)), Nat.add(Nat.mul(Nat.add(x, y), z), Nat.add(pa, pb)), Nat.add(Nat.add(xz, pa), Nat.add(yz, pb)),
54 Equal.cong(Nat, Nat, s => Nat.add(Nat.mul(Nat.add(x, y), z), s), Spec.dotn(Spec.vadd(xt, yt), zt), Nat.add(pa, pb), dot_vadd(xt, yt, zt)),
55 Equal.trans(Nat, Nat.add(Nat.mul(Nat.add(x, y), z), Nat.add(pa, pb)), Nat.add(Nat.add(xz, yz), Nat.add(pa, pb)), Nat.add(Nat.add(xz, pa), Nat.add(yz, pb)),
56 Equal.cong(Nat, Nat, s => Nat.add(s, Nat.add(pa, pb)), Nat.mul(Nat.add(x, y), z), Nat.add(xz, yz), Equal.sym(Nat, Nat.add(xz, yz), Nat.mul(Nat.add(x, y), z), NL.mul_dist(x, y, z))),
57 add4(xz, yz, pa, pb)))
58
59# dot(k·r, v) == k · dot(r, v): at each entry (k·x)·z == k·(x·z) (associativity), then k factors out of the sum
60def dot_scale(+k: Nat, r: List<&2, Nat>, v: List<&2, Nat>) -> {Spec.dotn(Spec.vscale(k, r), v) == Nat.mul(k, Spec.dotn(r, v)) : Nat}:
61 match r v:
62 case Nil{} _:
63 Equal.sym(Nat, Nat.mul(k, 0n), 0n, NL.mul_zero(k))
64 case Con{x, xt} Nil{}:
65 Equal.sym(Nat, Nat.mul(k, 0n), 0n, NL.mul_zero(k))
66 case Con{+x, +xt} Con{+z, +zt}:
67 +d = Spec.dotn(xt, zt)
68 Equal.trans(Nat, Nat.add(Nat.mul(Nat.mul(k, x), z), Spec.dotn(Spec.vscale(k, xt), zt)), Nat.add(Nat.mul(Nat.mul(k, x), z), Nat.mul(k, d)), Nat.mul(k, Nat.add(Nat.mul(x, z), d)),
69 Equal.cong(Nat, Nat, s => Nat.add(Nat.mul(Nat.mul(k, x), z), s), Spec.dotn(Spec.vscale(k, xt), zt), Nat.mul(k, d), dot_scale(k, xt, zt)),
70 Equal.trans(Nat, Nat.add(Nat.mul(Nat.mul(k, x), z), Nat.mul(k, d)), Nat.add(Nat.mul(k, Nat.mul(x, z)), Nat.mul(k, d)), Nat.mul(k, Nat.add(Nat.mul(x, z), d)),
71 Equal.cong(Nat, Nat, s => Nat.add(s, Nat.mul(k, d)), Nat.mul(Nat.mul(k, x), z), Nat.mul(k, Nat.mul(x, z)), Equal.sym(Nat, Nat.mul(k, Nat.mul(x, z)), Nat.mul(Nat.mul(k, x), z), NL.mul_assoc(k, x, z))),
72 Equal.sym(Nat, Nat.mul(k, Nat.add(Nat.mul(x, z), d)), Nat.add(Nat.mul(k, Nat.mul(x, z)), Nat.mul(k, d)), mul_dist_l(k, Nat.mul(x, z), d))))
73
74# dot(a, b) == dot(b, a)
75def dot_comm(a: List<&2, Nat>, b: List<&2, Nat>) -> {Spec.dotn(a, b) == Spec.dotn(b, a) : Nat}:
76 match a b:
77 case Nil{} Nil{}:
78 {==}
79 case Nil{} Con{y, yt}:
80 {==}
81 case Con{x, xt} Nil{}:
82 {==}
83 case Con{+x, +xt} Con{+y, +yt}:
84 Equal.trans(Nat, Nat.add(Nat.mul(x, y), Spec.dotn(xt, yt)), Nat.add(Nat.mul(y, x), Spec.dotn(xt, yt)), Nat.add(Nat.mul(y, x), Spec.dotn(yt, xt)),
85 Equal.cong(Nat, Nat, s => Nat.add(s, Spec.dotn(xt, yt)), Nat.mul(x, y), Nat.mul(y, x), NL.mul_comm(x, y)),
86 Equal.cong(Nat, Nat, s => Nat.add(Nat.mul(y, x), s), Spec.dotn(xt, yt), Spec.dotn(yt, xt), dot_comm(xt, yt)))
87
88# =====================================================================================================
89# The linear laws
90# =====================================================================================================
91
92# PROOF (adjoint): induction on the rows of W and the entries of dy. Row r with weight g_i contributes g_i·(r·x) on the
93# left and (g_i·r)·x on the right (dot_scale); the right side's sum of weighted rows splits by dot_vadd.
94def Laws.adjoint(w, g, v):
95 match w g:
96 case Nil{} Nil{}:
97 {==}
98 case Nil{} Con{gi, gt}:
99 {==}
100 case Con{r, rt} Nil{}:
101 {==}
102 case Con{+r, +rt} Con{+gi, +gt}:
103 +a = Spec.dotn(Spec.vscale(gi, r), v)
104 +m = Spec.mtv(rt, gt)
105 Equal.trans(Nat, Nat.add(Nat.mul(gi, Spec.dotn(r, v)), Spec.dotn(gt, Spec.mv(rt, v))), Nat.add(Nat.mul(gi, Spec.dotn(r, v)), Spec.dotn(m, v)), Spec.dotn(Spec.vadd(Spec.vscale(gi, r), m), v),
106 Equal.cong(Nat, Nat, s => Nat.add(Nat.mul(gi, Spec.dotn(r, v)), s), Spec.dotn(gt, Spec.mv(rt, v)), Spec.dotn(m, v), Laws.adjoint(rt, gt, v)),
107 Equal.trans(Nat, Nat.add(Nat.mul(gi, Spec.dotn(r, v)), Spec.dotn(m, v)), Nat.add(a, Spec.dotn(m, v)), Spec.dotn(Spec.vadd(Spec.vscale(gi, r), m), v),
108 Equal.cong(Nat, Nat, s => Nat.add(s, Spec.dotn(m, v)), Nat.mul(gi, Spec.dotn(r, v)), a, Equal.sym(Nat, a, Nat.mul(gi, Spec.dotn(r, v)), dot_scale(gi, r, v))),
109 Equal.sym(Nat, Spec.dotn(Spec.vadd(Spec.vscale(gi, r), m), v), Nat.add(a, Spec.dotn(m, v)), dot_vadd(Spec.vscale(gi, r), m, v))))
110
111# PROOF (weight_adjoint): row by row, g_i·(r·x) == (g_i·x)·r, by dot_comm and dot_scale.
112def Laws.weight_adjoint(w, g, v):
113 match w g:
114 case Nil{} Nil{}:
115 {==}
116 case Nil{} Con{gi, gt}:
117 {==}
118 case Con{r, rt} Nil{}:
119 {==}
120 case Con{+r, +rt} Con{+gi, +gt}:
121 +f = Spec.frob(Spec.outer(gt, v), rt)
122 Equal.trans(Nat, Nat.add(Nat.mul(gi, Spec.dotn(r, v)), Spec.dotn(gt, Spec.mv(rt, v))), Nat.add(Nat.mul(gi, Spec.dotn(r, v)), f), Nat.add(Spec.dotn(Spec.vscale(gi, v), r), f),
123 Equal.cong(Nat, Nat, s => Nat.add(Nat.mul(gi, Spec.dotn(r, v)), s), Spec.dotn(gt, Spec.mv(rt, v)), f, Laws.weight_adjoint(rt, gt, v)),
124 Equal.cong(Nat, Nat, s => Nat.add(s, f), Nat.mul(gi, Spec.dotn(r, v)), Spec.dotn(Spec.vscale(gi, v), r),
125 Equal.trans(Nat, Nat.mul(gi, Spec.dotn(r, v)), Nat.mul(gi, Spec.dotn(v, r)), Spec.dotn(Spec.vscale(gi, v), r),
126 Equal.cong(Nat, Nat, s => Nat.mul(gi, s), Spec.dotn(r, v), Spec.dotn(v, r), dot_comm(r, v)),
127 Equal.sym(Nat, Spec.dotn(Spec.vscale(gi, v), r), Nat.mul(gi, Spec.dotn(v, r)), dot_scale(gi, v, r)))))
128
129# PROOF (adjoint_seq): adjoint once for W2 (with W1·x as the input), then once for W1.
130def Laws.adjoint_seq(w1, w2, g, v):
131 Equal.trans(Nat, Spec.dotn(g, Spec.mv(w2, Spec.mv(w1, v))), Spec.dotn(Spec.mtv(w2, g), Spec.mv(w1, v)), Spec.dotn(Spec.mtv(w1, Spec.mtv(w2, g)), v), Laws.adjoint(w2, g, Spec.mv(w1, v)), Laws.adjoint(w1, Spec.mtv(w2, g), v))
132
133# =====================================================================================================
134# Causal attention
135# =====================================================================================================
136
137# PROOF (seq_chain): Seq's forward pass computes W2·(W1·x) and saves x and W1·x; its backward pass computes
138# W1ᵀ·(W2ᵀ·dy). Every step of the combinator reduces, so this is adjoint_seq.
139def Laws.seq_chain(w1, w2, x, dy, g1, g2):
140 Laws.adjoint_seq(w1, w2, dy, x)
141
142# PROOF (fork_sum): Fork's backward pass is W1ᵀ·dy1 + W2ᵀ·dy2 (vadd, the add it was given); the dot of a sum is the sum
143# of the dots (dot_vadd), and each is its layer's adjoint law.
144def Laws.fork_sum(w1, w2, x, dy1, dy2, g1, g2):
145 Equal.trans(Nat, Nat.add(Spec.dotn(dy1, Spec.mv(w1, x)), Spec.dotn(dy2, Spec.mv(w2, x))), Nat.add(Spec.dotn(Spec.mtv(w1, dy1), x), Spec.dotn(Spec.mtv(w2, dy2), x)), Spec.dotn(Spec.vadd(Spec.mtv(w1, dy1), Spec.mtv(w2, dy2)), x),
146 Equal.trans(Nat, Nat.add(Spec.dotn(dy1, Spec.mv(w1, x)), Spec.dotn(dy2, Spec.mv(w2, x))), Nat.add(Spec.dotn(Spec.mtv(w1, dy1), x), Spec.dotn(dy2, Spec.mv(w2, x))), Nat.add(Spec.dotn(Spec.mtv(w1, dy1), x), Spec.dotn(Spec.mtv(w2, dy2), x)),
147 Equal.cong(Nat, Nat, z => Nat.add(z, Spec.dotn(dy2, Spec.mv(w2, x))), Spec.dotn(dy1, Spec.mv(w1, x)), Spec.dotn(Spec.mtv(w1, dy1), x), Laws.adjoint(w1, dy1, x)),
148 Equal.cong(Nat, Nat, z => Nat.add(Spec.dotn(Spec.mtv(w1, dy1), x), z), Spec.dotn(dy2, Spec.mv(w2, x)), Spec.dotn(Spec.mtv(w2, dy2), x), Laws.adjoint(w2, dy2, x))),
149 Equal.sym(Nat, Spec.dotn(Spec.vadd(Spec.mtv(w1, dy1), Spec.mtv(w2, dy2)), x), Nat.add(Spec.dotn(Spec.mtv(w1, dy1), x), Spec.dotn(Spec.mtv(w2, dy2), x)), dot_vadd(Spec.mtv(w1, dy1), Spec.mtv(w2, dy2), x)))
150
151def disc(b: Bool) -> Type:
152 match b:
153 case True{}:
154 Unit
155 case False{}:
156 Empty
157
158# False == True is impossible: rewriting the goal Empty = disc(False) with it gives disc(True) = Unit
159def false_true(e: {False{} == True{} : Bool}) -> Empty:
160 %Equal.sym(Bool, False{}, True{}, e) : disc(_)
161 Unit{}
162
163# xs reversed onto acc has the length of both
164def len_rcat(+xs: List<&2, F32>, +acc: List<&2, F32>) -> {List.length(&2, F32, T.rcat(xs, acc)) == Nat.add(List.length(&2, F32, xs), List.length(&2, F32, acc)) : Nat}:
165 match xs:
166 case Nil{}:
167 {==}
168 case Con{+x, +t}:
169 Equal.trans(Nat, List.length(&2, F32, T.rcat(t, x <> acc)), Nat.add(List.length(&2, F32, t), 1n+List.length(&2, F32, acc)), 1n+Nat.add(List.length(&2, F32, t), List.length(&2, F32, acc)), len_rcat(t, x <> acc), Equal.sym(Nat, 1n+Nat.add(List.length(&2, F32, t), List.length(&2, F32, acc)), Nat.add(List.length(&2, F32, t), 1n+List.length(&2, F32, acc)), NL.add_succ(List.length(&2, F32, t), List.length(&2, F32, acc))))
170
171# drow's row has the current dot, the m after it, and acc: 1 + m + |acc| numbers
172def len_drow(m: Nat, +hd: Nat, +ia: U32, +ib: U32, +bs: U32, +acc: List<&2, F32>, r: T.D2) -> {List.length(&2, F32, T.dr_row(T.drow(m, hd, ia, ib, bs, acc, r))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat}:
173 match m r:
174 case 0n T.D2{a, b, +s}:
175 Equal.trans(Nat, List.length(&2, F32, T.rcat(s <> acc, Nil{})), Nat.add(1n+List.length(&2, F32, acc), 0n), 1n+List.length(&2, F32, acc), len_rcat(s <> acc, Nil{}), NL.add_zero(1n+List.length(&2, F32, acc)))
176 case 1n++p T.D2{a, b, +s}:
177 Equal.trans(Nat, List.length(&2, F32, T.dr_row(T.drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, T.dot0(hd, ia, (ib + bs : U32), 1, 1, a, b)))), Nat.add(1n+p, 1n+List.length(&2, F32, acc)), 2n+Nat.add(p, List.length(&2, F32, acc)), len_drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, T.dot0(hd, ia, (ib + bs : U32), 1, 1, a, b)), Equal.cong(Nat, Nat, z => 1n+z, Nat.add(p, 1n+List.length(&2, F32, acc)), 1n+Nat.add(p, List.length(&2, F32, acc)), Equal.sym(Nat, 1n+Nat.add(p, List.length(&2, F32, acc)), Nat.add(p, 1n+List.length(&2, F32, acc)), NL.add_succ(p, List.length(&2, F32, acc)))))
178
179# PROOF (causal_keys): qrow is drow(i, ..) from an empty accumulator, so its row has 1 + i + 0 numbers (len_drow).
180def Laws.causal_keys(i, hd, qs, qc, ks, kc, q, k):
181 Equal.trans(Nat, List.length(&2, F32, T.dr_row(T.qrow(i, hd, qs, qc, ks, kc, q, k))), Nat.add(1n+i, 0n), 1n+i, len_drow(i, hd, ((U32.from_nat(i) * qs : U32) + qc : U32), kc, ks, Nil{}, T.dot0(hd, ((U32.from_nat(i) * qs : U32) + qc : U32), kc, 1, 1, q, k)), NL.add_zero(1n+i))
182
183# n - 0 is n
184def sub0(+n: Nat) -> {Nat.sub(n, 0n) == n : Nat}:
185 match n:
186 case 0n:
187 {==}
188 case 1n+p:
189 {==}
190
191# the division loop of a softmax counts the 1 + k numbers it divides onto c (pr_of hands the count on as the row's)
192def cnt_div(q: Array<F32>, k: Array<F32>, kk: Nat, +idx: U32, +s: F32, +c: Nat, r: Array<F32> & F32) -> {T.pr_n(T.pr_of(q, k, T.pdiv(kk, idx, s, c, r))) == Nat.add(1n+kk, c) : Nat}:
193 match kk r:
194 case 0n Tuple{a, +x}:
195 {==}
196 case 1n++j Tuple{a, +x}:
197 Equal.trans(Nat, T.pr_n(T.pr_of(q, k, T.pdiv(j, (idx + 1 : U32), s, 1n+c, Array.get(F32, Array.set(F32, a, idx, (x / s : F32)), (idx + 1 : U32))))), Nat.add(1n+j, 1n+c), 2n+Nat.add(j, c), cnt_div(q, k, j, (idx + 1 : U32), s, 1n+c, Array.get(F32, Array.set(F32, a, idx, (x / s : F32)), (idx + 1 : U32))), Equal.cong(Nat, Nat, z => 1n+z, Nat.add(j, 1n+c), 1n+Nat.add(j, c), Equal.sym(Nat, 1n+Nat.add(j, c), Nat.add(j, 1n+c), NL.add_succ(j, c))))
198
199# the whole softmax over 1 + kk numbers reports 1 + kk, whatever max and exp found (they are matched, not computed)
200def cnt_psm3(q: Array<F32>, k: Array<F32>, +kk: Nat, +off: U32, r: T.AS) -> {T.pr_n(T.pr_of(q, k, T.psm3(kk, off, r))) == 1n+kk : Nat}:
201 match r:
202 case T.AS{a, +s}:
203 Equal.trans(Nat, T.pr_n(T.pr_of(q, k, T.pdiv(kk, off, s, 0n, Array.get(F32, a, off)))), Nat.add(1n+kk, 0n), 1n+kk, cnt_div(q, k, kk, off, s, 0n, Array.get(F32, a, off)), NL.add_zero(1n+kk))
204
205def cnt_psm2(q: Array<F32>, k: Array<F32>, +kk: Nat, +off: U32, r: T.AS) -> {T.pr_n(T.pr_of(q, k, T.psm2(kk, off, r))) == 1n+kk : Nat}:
206 match r:
207 case T.AS{a, +m}:
208 cnt_psm3(q, k, kk, off, T.pexp(kk, off, m, 0.0, Array.get(F32, a, off)))
209
210# pw writes xs and counts them from n; the softmax prow3 runs on them reports 1 + (|xs| + n - 1)
211def cnt_pw(+xs: List<&2, F32>, +idx: U32, +sc: F32, +n: Nat, +off: U32, q: Array<F32>, k: Array<F32>, p: Array<F32>) -> {T.pr_n(T.prow3(off, q, k, T.pw(xs, idx, sc, n, p))) == 1n+Nat.sub(Nat.add(List.length(&2, F32, xs), n), 1n) : Nat}:
212 match xs:
213 case Nil{}:
214 cnt_psm2(q, k, Nat.sub(n, 1n), off, T.pmax0(Nat.sub(n, 1n), off, Array.get(F32, p, off)))
215 case Con{+x, +t}:
216 Equal.trans(Nat, T.pr_n(T.prow3(off, q, k, T.pw(t, (idx + 1 : U32), sc, 1n+n, Array.set(F32, p, idx, (x * sc : F32))))), 1n+Nat.sub(Nat.add(List.length(&2, F32, t), 1n+n), 1n), 1n+Nat.sub(1n+Nat.add(List.length(&2, F32, t), n), 1n), cnt_pw(t, (idx + 1 : U32), sc, 1n+n, off, q, k, Array.set(F32, p, idx, (x * sc : F32))), Equal.cong(Nat, Nat, z => 1n+Nat.sub(z, 1n), Nat.add(List.length(&2, F32, t), 1n+n), 1n+Nat.add(List.length(&2, F32, t), n), Equal.sym(Nat, 1n+Nat.add(List.length(&2, F32, t), n), Nat.add(List.length(&2, F32, t), 1n+n), NL.add_succ(List.length(&2, F32, t), n))))
217
218# the count prow2 hands on for drow's row: 1 + m + |acc|, as len_drow, through cnt_pw
219def cnt_drow(m: Nat, +hd: Nat, +ia: U32, +ib: U32, +bs: U32, +acc: List<&2, F32>, r: T.D2, +sc: F32, +off: U32, pa: Array<F32>) -> {T.pr_n(T.prow2(sc, off, pa, T.drow(m, hd, ia, ib, bs, acc, r))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat}:
220 match m r:
221 case 0n T.D2{a, b, +s}:
222 +rw = T.rcat(s <> acc, Nil{})
223 +la = List.length(&2, F32, acc)
224 Equal.trans(Nat, T.pr_n(T.prow3(off, a, b, T.pw(rw, off, sc, 0n, pa))), 1n+Nat.sub(Nat.add(List.length(&2, F32, rw), 0n), 1n), 1n+la, cnt_pw(rw, off, sc, 0n, off, a, b, pa),
225 Equal.cong(Nat, Nat, z => 1n+z, Nat.sub(Nat.add(List.length(&2, F32, rw), 0n), 1n), la,
226 Equal.trans(Nat, Nat.sub(Nat.add(List.length(&2, F32, rw), 0n), 1n), Nat.sub(1n+la, 1n), la,
227 Equal.cong(Nat, Nat, z => Nat.sub(z, 1n), Nat.add(List.length(&2, F32, rw), 0n), 1n+la,
228 Equal.trans(Nat, Nat.add(List.length(&2, F32, rw), 0n), List.length(&2, F32, rw), 1n+la, NL.add_zero(List.length(&2, F32, rw)),
229 Equal.trans(Nat, List.length(&2, F32, rw), Nat.add(1n+la, 0n), 1n+la, len_rcat(s <> acc, Nil{}), NL.add_zero(1n+la)))),
230 sub0(la))))
231 case 1n++p T.D2{a, b, +s}:
232 Equal.trans(Nat, T.pr_n(T.prow2(sc, off, pa, T.drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, T.dot0(hd, ia, (ib + bs : U32), 1, 1, a, b)))), Nat.add(1n+p, 1n+List.length(&2, F32, acc)), 2n+Nat.add(p, List.length(&2, F32, acc)), cnt_drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, T.dot0(hd, ia, (ib + bs : U32), 1, 1, a, b), sc, off, pa), Equal.cong(Nat, Nat, z => 1n+z, Nat.add(p, 1n+List.length(&2, F32, acc)), 1n+Nat.add(p, List.length(&2, F32, acc)), Equal.sym(Nat, 1n+Nat.add(p, List.length(&2, F32, acc)), Nat.add(p, 1n+List.length(&2, F32, acc)), NL.add_succ(p, List.length(&2, F32, acc)))))
233
234# PROOF (causal): prow runs prow2 on qrow's row, drow(i, ..) from an empty accumulator, so the count it hands the
235# softmax is 1 + i + 0 (cnt_drow).
236def Laws.causal(i, t, hd, sc, qs, qc, ks, kc, q, k, p):
237 Equal.trans(Nat, T.pr_n(T.prow(i, t, hd, sc, qs, qc, ks, kc, q, k, p)), Nat.add(1n+i, 0n), 1n+i, cnt_drow(i, hd, ((U32.from_nat(i) * qs : U32) + qc : U32), kc, ks, Nil{}, T.dot0(hd, ((U32.from_nat(i) * qs : U32) + qc : U32), kc, 1, 1, q, k), sc, (U32.from_nat(i) * U32.from_nat(t) : U32), p), NL.add_zero(1n+i))
238
239# gw writes xs and counts them from n: gr3 hands on |xs| + n (as cnt_pw for the forward pass)
240def cnt_gw(+xs: List<&2, F32>, +idx: U32, +n: Nat, +sc: F32, +off: U32, d: Array<F32>, v: Array<F32>, p: Array<F32>, g: Array<F32>) -> {T.gr_n(T.gr3(sc, off, d, v, p, T.gw(xs, idx, n, g))) == Nat.add(List.length(&2, F32, xs), n) : Nat}:
241 match xs:
242 case Nil{}:
243 {==}
244 case Con{+x, +t}:
245 Equal.trans(Nat, T.gr_n(T.gr3(sc, off, d, v, p, T.gw(t, (idx + 1 : U32), 1n+n, Array.set(F32, g, idx, x)))), Nat.add(List.length(&2, F32, t), 1n+n), 1n+Nat.add(List.length(&2, F32, t), n), cnt_gw(t, (idx + 1 : U32), 1n+n, sc, off, d, v, p, Array.set(F32, g, idx, x)), Equal.sym(Nat, 1n+Nat.add(List.length(&2, F32, t), n), Nat.add(List.length(&2, F32, t), 1n+n), NL.add_succ(List.length(&2, F32, t), n)))
246
247# the count gr2 hands on for drow's row: 1 + m + |acc|, as len_drow, through cnt_gw (cnt_drow's twin)
248def cnt_gdrow(m: Nat, +hd: Nat, +ia: U32, +ib: U32, +bs: U32, +acc: List<&2, F32>, r: T.D2, +sc: F32, +off: U32, pa: Array<F32>, ga: Array<F32>) -> {T.gr_n(T.gr2(sc, off, pa, ga, T.drow(m, hd, ia, ib, bs, acc, r))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat}:
249 match m r:
250 case 0n T.D2{a, b, +s}:
251 +rw = T.rcat(s <> acc, Nil{})
252 Equal.trans(Nat, T.gr_n(T.gr3(sc, off, a, b, pa, T.gw(rw, off, 0n, ga))), Nat.add(List.length(&2, F32, rw), 0n), 1n+List.length(&2, F32, acc), cnt_gw(rw, off, 0n, sc, off, a, b, pa, ga),
253 Equal.trans(Nat, Nat.add(List.length(&2, F32, rw), 0n), List.length(&2, F32, rw), 1n+List.length(&2, F32, acc), NL.add_zero(List.length(&2, F32, rw)),
254 Equal.trans(Nat, List.length(&2, F32, rw), Nat.add(1n+List.length(&2, F32, acc), 0n), 1n+List.length(&2, F32, acc), len_rcat(s <> acc, Nil{}), NL.add_zero(1n+List.length(&2, F32, acc)))))
255 case 1n++p T.D2{a, b, +s}:
256 Equal.trans(Nat, T.gr_n(T.gr2(sc, off, pa, ga, T.drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, T.dot0(hd, ia, (ib + bs : U32), 1, 1, a, b)))), Nat.add(1n+p, 1n+List.length(&2, F32, acc)), 2n+Nat.add(p, List.length(&2, F32, acc)), cnt_gdrow(p, hd, ia, (ib + bs : U32), bs, s <> acc, T.dot0(hd, ia, (ib + bs : U32), 1, 1, a, b), sc, off, pa, ga), Equal.cong(Nat, Nat, z => 1n+z, Nat.add(p, 1n+List.length(&2, F32, acc)), 1n+Nat.add(p, List.length(&2, F32, acc)), Equal.sym(Nat, 1n+Nat.add(p, List.length(&2, F32, acc)), Nat.add(p, 1n+List.length(&2, F32, acc)), NL.add_succ(p, List.length(&2, F32, acc)))))
257
258# PROOF (causal_grad): grow runs gr2 on drow(i, ..) from an empty accumulator, so the count is 1 + i + 0 (cnt_gdrow).
259def Laws.causal_grad(i, t, hd, sc, ds, dc, vs, vc, d, v, p, g):
260 Equal.trans(Nat, T.gr_n(T.grow(i, t, hd, sc, ds, dc, vs, vc, d, v, p, g)), Nat.add(1n+i, 0n), 1n+i, cnt_gdrow(i, hd, ((U32.from_nat(i) * ds : U32) + dc : U32), vc, vs, Nil{}, T.dot0(hd, ((U32.from_nat(i) * ds : U32) + dc : U32), vc, 1, 1, d, v), sc, (U32.from_nat(i) * U32.from_nat(t) : U32), p, g), NL.add_zero(1n+i))
261
262# =====================================================================================================
263# Attention's trees of sequences: parallel = sequential
264# =====================================================================================================
265
266# a ++ b == a2 ++ b2 from a == a2 and b == b2
267def app2(-A: Type, -a: List<&1, A>, -a2: List<&1, A>, -b: List<&1, A>, -b2: List<&1, A>, ea: {a == a2 : List<&1, A>}, eb: {b == b2 : List<&1, A>}) -> {List.append(&1, A, a, b) == List.append(&1, A, a2, b2) : List<&1, A>}:
268 Equal.trans(List<&1, A>, List.append(&1, A, a, b), List.append(&1, A, a2, b), List.append(&1, A, a2, b2), Equal.cong(List<&1, A>, List<&1, A>, z => List.append(&1, A, z, b), a, a2, ea), Equal.cong(List<&1, A>, List<&1, A>, z => List.append(&1, A, a2, z), b, b2, eb))
269
270# ---- the sequential map the tree is compared with: f(x) = fs_o(fheads(h, 0, dm, x)) for each x, in order ----
271
272def fmap(+h: Nat, +dm: T.Dm, xs: List<&1, T.FS>) -> List<&1, Array<F32>>:
273 match xs:
274 case Nil{}:
275 Nil{}
276 case Con{x, t}:
277 T.fs_o(T.fheads(h, 0n, dm, x)) <> fmap(h, dm, t)
278
279# ---- fleaf (the code's loop with an accumulator) is the map pushed onto the accumulator ----
280
281type L1R<-h: Nat, -dm: T.Dm, -xs: List<&1, T.FS>, -acc: List<&1, Array<F32>>> is Type:
282 L1R{ys: List<&1, T.FS>, same: {ys == xs : List<&1, T.FS>}, fact: {T.fleaf(xs, h, dm, acc) == List.reverse.go(&1, Array<F32>, acc, fmap(h, dm, xs)) : List<&1, Array<F32>>}}
283
284def l1_con(x: T.FS, -t: List<&1, T.FS>, +h: Nat, +dm: T.Dm, -acc: List<&1, Array<F32>>, r: L1R<h, dm, t, T.fs_o(T.fheads(h, 0n, dm, x)) <> acc>) -> L1R<h, dm, x <> t, acc>:
285 match r:
286 case L1R{ys, same, fact}:
287 L1R{x <> ys, Equal.cong(List<&1, T.FS>, List<&1, T.FS>, z => x <> z, ys, t, same), fact}
288
289def l1(xs: List<&1, T.FS>, +h: Nat, +dm: T.Dm, -acc: List<&1, Array<F32>>) -> L1R<h, dm, xs, acc>:
290 match xs:
291 case Nil{}:
292 L1R{Nil{}, {==}, {==}}
293 case Con{x, t}:
294 l1_con(x, t, h, dm, acc, l1(t, h, dm, T.fs_o(T.fheads(h, 0n, dm, x)) <> acc))
295
296# ---- the map of two lists appended is the two maps appended ----
297
298type MAR<-h: Nat, -dm: T.Dm, -a: List<&1, T.FS>, -b: List<&1, T.FS>> is Type:
299 MAR{ys: List<&1, T.FS>, same: {ys == a : List<&1, T.FS>}, fact: {fmap(h, dm, List.append(&1, T.FS, a, b)) == List.append(&1, Array<F32>, fmap(h, dm, a), fmap(h, dm, b)) : List<&1, Array<F32>>}}
300
301def ma_con(x: T.FS, -t: List<&1, T.FS>, +h: Nat, +dm: T.Dm, -b: List<&1, T.FS>, r: MAR<h, dm, t, b>) -> MAR<h, dm, x <> t, b>:
302 match r:
303 case MAR{ys, same, fact}:
304 MAR{x <> ys, Equal.cong(List<&1, T.FS>, List<&1, T.FS>, z => x <> z, ys, t, same), Equal.cong(List<&1, Array<F32>>, List<&1, Array<F32>>, z => T.fs_o(T.fheads(h, 0n, dm, x)) <> z, fmap(h, dm, List.append(&1, T.FS, t, b)), List.append(&1, Array<F32>, fmap(h, dm, t), fmap(h, dm, b)), fact)}
305
306def map_app(a: List<&1, T.FS>, +h: Nat, +dm: T.Dm, -b: List<&1, T.FS>) -> MAR<h, dm, a, b>:
307 match a:
308 case Nil{}:
309 MAR{Nil{}, {==}, {==}}
310 case Con{x, t}:
311 ma_con(x, t, h, dm, b, map_app(t, h, dm, b))
312
313# ---- a reversed onto b, then c appended: a reversed onto b ++ c ----
314
315type GAR<-a: List<&1, T.FS>, -b: List<&1, T.FS>, -c: List<&1, T.FS>> is Type:
316 GAR{ys: List<&1, T.FS>, same: {ys == a : List<&1, T.FS>}, fact: {List.append(&1, T.FS, List.reverse.go(&1, T.FS, a, b), c) == List.reverse.go(&1, T.FS, a, List.append(&1, T.FS, b, c)) : List<&1, T.FS>}}
317
318def ga_con(x: T.FS, -t: List<&1, T.FS>, -b: List<&1, T.FS>, -c: List<&1, T.FS>, r: GAR<t, x <> b, c>) -> GAR<x <> t, b, c>:
319 match r:
320 case GAR{ys, same, fact}:
321 GAR{x <> ys, Equal.cong(List<&1, T.FS>, List<&1, T.FS>, z => x <> z, ys, t, same), fact}
322
323def go_app(a: List<&1, T.FS>, -b: List<&1, T.FS>, -c: List<&1, T.FS>) -> GAR<a, b, c>:
324 match a:
325 case Nil{}:
326 GAR{Nil{}, {==}, {==}}
327 case Con{x, t}:
328 ga_con(x, t, b, c, go_app(t, x <> b, c))
329
330# ---- cut gives two halves that, appended, are acc reversed onto xs ----
331
332type CGR<-k: Nat, -xs: List<&1, T.FS>, -acc: List<&1, T.FS>> is Type:
333 CGR{hh: T.Halves<T.FS>, same: {hh == T.cut(T.FS, k, xs, acc) : T.Halves<T.FS>}, cat: {T.hcat(T.FS, hh) == List.reverse.go(&1, T.FS, acc, xs) : List<&1, T.FS>}}
334
335def cg_base(-k: Nat, xs: List<&1, T.FS>, -acc: List<&1, T.FS>, e: {T.cut(T.FS, k, xs, acc) == T.Halves{List.reverse(&1, T.FS, acc), xs} : T.Halves<T.FS>}, r: GAR<acc, Nil{}, xs>) -> CGR<k, xs, acc>:
336 match r:
337 case GAR{ys, +same, fact}:
338 CGR{T.Halves{List.reverse(&1, T.FS, ys), xs},
339 Equal.trans(T.Halves<T.FS>, T.Halves{List.reverse(&1, T.FS, ys), xs}, T.Halves{List.reverse(&1, T.FS, acc), xs}, T.cut(T.FS, k, xs, acc), Equal.cong(List<&1, T.FS>, T.Halves<T.FS>, z => T.Halves{List.reverse(&1, T.FS, z), xs}, ys, acc, same), Equal.sym(T.Halves<T.FS>, T.cut(T.FS, k, xs, acc), T.Halves{List.reverse(&1, T.FS, acc), xs}, e)),
340 Equal.trans(List<&1, T.FS>, List.append(&1, T.FS, List.reverse.go(&1, T.FS, ys, Nil{}), xs), List.append(&1, T.FS, List.reverse.go(&1, T.FS, acc, Nil{}), xs), List.reverse.go(&1, T.FS, acc, xs), Equal.cong(List<&1, T.FS>, List<&1, T.FS>, z => List.append(&1, T.FS, List.reverse.go(&1, T.FS, z, Nil{}), xs), ys, acc, same), fact)}
341
342def cg_con(-m: Nat, -x: T.FS, -t: List<&1, T.FS>, -acc: List<&1, T.FS>, r: CGR<m, t, x <> acc>) -> CGR<1n+m, x <> t, acc>:
343 match r:
344 case CGR{hh, same, cat}:
345 CGR{hh, same, cat}
346
347def cg(k: Nat, xs: List<&1, T.FS>, acc: List<&1, T.FS>) -> CGR<k, xs, acc>:
348 match k xs:
349 case 1n+m Con{x, t}:
350 cg_con(m, x, t, acc, cg(m, t, x <> acc))
351 case 0n ys:
352 cg_base(0n, ys, acc, {==}, go_app(acc, Nil{}, ys))
353 case 1n+m Nil{}:
354 cg_base(1n+m, Nil{}, acc, {==}, go_app(acc, Nil{}, Nil{}))
355
356# ---- the tree: fpar at depth d on two halves is the map of the halves appended ----
357
358type MR<-d: Nat, -gs: T.Halves<T.FS>, -n: Nat, -h: Nat, -dm: T.Dm> is Type:
359 MR{ys: List<&1, T.FS>, same: {ys == T.hcat(T.FS, gs) : List<&1, T.FS>}, fact: {T.fpar(d, gs, n, h, dm) == fmap(h, dm, T.hcat(T.FS, gs)) : List<&1, Array<F32>>}}
360
361# depth 0: the two leaves' loops, appended
362def m0c(-l: List<&1, T.FS>, -r: List<&1, T.FS>, -n: Nat, +h: Nat, +dm: T.Dm, -l1y: List<&1, T.FS>, +sl: {l1y == l : List<&1, T.FS>}, fl: {T.fleaf(l, h, dm, Nil{}) == fmap(h, dm, l) : List<&1, Array<F32>>}, r1: List<&1, T.FS>, sr: {r1 == r : List<&1, T.FS>}, fr: {T.fleaf(r, h, dm, Nil{}) == fmap(h, dm, r) : List<&1, Array<F32>>}, c: MAR<h, dm, l1y, r>) -> MR<0n, T.Halves{l, r}, n, h, dm>:
363 match c:
364 case MAR{l2, s2, fa}:
365 MR{List.append(&1, T.FS, l2, r1), app2(T.FS, l2, l, r1, r, Equal.trans(List<&1, T.FS>, l2, l1y, l, s2, sl), sr),
366 Equal.trans(List<&1, Array<F32>>, List.append(&1, Array<F32>, T.fleaf(l, h, dm, Nil{}), T.fleaf(r, h, dm, Nil{})), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), fmap(h, dm, List.append(&1, T.FS, l, r)), app2(Array<F32>, T.fleaf(l, h, dm, Nil{}), fmap(h, dm, l), T.fleaf(r, h, dm, Nil{}), fmap(h, dm, r), fl, fr),
367 Equal.sym(List<&1, Array<F32>>, fmap(h, dm, List.append(&1, T.FS, l, r)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)),
368 Equal.trans(List<&1, Array<F32>>, fmap(h, dm, List.append(&1, T.FS, l, r)), fmap(h, dm, List.append(&1, T.FS, l1y, r)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, List.append(&1, T.FS, z, r)), l, l1y, Equal.sym(List<&1, T.FS>, l1y, l, sl)),
369 Equal.trans(List<&1, Array<F32>>, fmap(h, dm, List.append(&1, T.FS, l1y, r)), List.append(&1, Array<F32>, fmap(h, dm, l1y), fmap(h, dm, r)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), fa, Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => List.append(&1, Array<F32>, fmap(h, dm, z), fmap(h, dm, r)), l1y, l, sl)))))}
370
371def m0b(-l: List<&1, T.FS>, -r: List<&1, T.FS>, -n: Nat, +h: Nat, +dm: T.Dm, l1y: List<&1, T.FS>, +sl: {l1y == l : List<&1, T.FS>}, fl: {T.fleaf(l, h, dm, Nil{}) == fmap(h, dm, l) : List<&1, Array<F32>>}, b: L1R<h, dm, r, Nil{}>) -> MR<0n, T.Halves{l, r}, n, h, dm>:
372 match b:
373 case L1R{r1, sr, fr}:
374 m0c(l, r, n, h, dm, l1y, sl, fl, r1, sr, fr, map_app(l1y, h, dm, r))
375
376def m0a(-l: List<&1, T.FS>, r: List<&1, T.FS>, -n: Nat, +h: Nat, +dm: T.Dm, a: L1R<h, dm, l, Nil{}>) -> MR<0n, T.Halves{l, r}, n, h, dm>:
377 match a:
378 case L1R{l1y, sl, fl}:
379 m0b(l, r, n, h, dm, l1y, sl, fl, l1(r, h, dm, Nil{}))
380
381# depth 1 + e: each half cut in two and its tree (the closures run main at depth e), then the two maps appended
382def m1c(-e: Nat, -l: List<&1, T.FS>, -r: List<&1, T.FS>, +n: Nat, +h: Nat, +dm: T.Dm, -hh1: T.Halves<T.FS>, s1: {hh1 == T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}) : T.Halves<T.FS>}, +cat1: {T.hcat(T.FS, hh1) == l : List<&1, T.FS>}, -hh2: T.Halves<T.FS>, s2: {hh2 == T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}) : T.Halves<T.FS>}, +cat2: {T.hcat(T.FS, hh2) == r : List<&1, T.FS>}, -y1: List<&1, T.FS>, sy1: {y1 == T.hcat(T.FS, hh1) : List<&1, T.FS>}, f1: {T.fpar(e, hh1, TA.half(n), h, dm) == fmap(h, dm, T.hcat(T.FS, hh1)) : List<&1, Array<F32>>}, y2: List<&1, T.FS>, sy2: {y2 == T.hcat(T.FS, hh2) : List<&1, T.FS>}, f2: {T.fpar(e, hh2, Nat.sub(n, TA.half(n)), h, dm) == fmap(h, dm, T.hcat(T.FS, hh2)) : List<&1, Array<F32>>}, ma: MAR<h, dm, y1, y2>) -> MR<1n+e, T.Halves{l, r}, n, h, dm>:
383 match ma:
384 case MAR{y1b, sb, fa}:
385 +ly1 = Equal.trans(List<&1, T.FS>, y1, T.hcat(T.FS, hh1), l, sy1, cat1)
386 +ly2 = Equal.trans(List<&1, T.FS>, y2, T.hcat(T.FS, hh2), r, sy2, cat2)
387 +q1 = Equal.trans(List<&1, Array<F32>>, T.fpar(e, T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm), T.fpar(e, hh1, TA.half(n), h, dm), fmap(h, dm, l), Equal.cong(T.Halves<T.FS>, List<&1, Array<F32>>, z => T.fpar(e, z, TA.half(n), h, dm), T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}), hh1, Equal.sym(T.Halves<T.FS>, hh1, T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}), s1)), Equal.trans(List<&1, Array<F32>>, T.fpar(e, hh1, TA.half(n), h, dm), fmap(h, dm, T.hcat(T.FS, hh1)), fmap(h, dm, l), f1, Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, z), T.hcat(T.FS, hh1), l, cat1)))
388 +q2 = Equal.trans(List<&1, Array<F32>>, T.fpar(e, T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm), T.fpar(e, hh2, Nat.sub(n, TA.half(n)), h, dm), fmap(h, dm, r), Equal.cong(T.Halves<T.FS>, List<&1, Array<F32>>, z => T.fpar(e, z, Nat.sub(n, TA.half(n)), h, dm), T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), hh2, Equal.sym(T.Halves<T.FS>, hh2, T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), s2)), Equal.trans(List<&1, Array<F32>>, T.fpar(e, hh2, Nat.sub(n, TA.half(n)), h, dm), fmap(h, dm, T.hcat(T.FS, hh2)), fmap(h, dm, r), f2, Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, z), T.hcat(T.FS, hh2), r, cat2)))
389 +bb = Equal.trans(List<&1, Array<F32>>, fmap(h, dm, List.append(&1, T.FS, l, r)), fmap(h, dm, List.append(&1, T.FS, y1, y2)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, z), List.append(&1, T.FS, l, r), List.append(&1, T.FS, y1, y2), app2(T.FS, l, y1, r, y2, Equal.sym(List<&1, T.FS>, y1, l, ly1), Equal.sym(List<&1, T.FS>, y2, r, ly2))), Equal.trans(List<&1, Array<F32>>, fmap(h, dm, List.append(&1, T.FS, y1, y2)), List.append(&1, Array<F32>, fmap(h, dm, y1), fmap(h, dm, y2)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), fa, app2(Array<F32>, fmap(h, dm, y1), fmap(h, dm, l), fmap(h, dm, y2), fmap(h, dm, r), Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, z), y1, l, ly1), Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, z), y2, r, ly2))))
390 MR{List.append(&1, T.FS, y1b, y2), app2(T.FS, y1b, l, y2, r, Equal.trans(List<&1, T.FS>, y1b, y1, l, sb, ly1), ly2),
391 Equal.trans(List<&1, Array<F32>>, List.append(&1, Array<F32>, T.fpar(e, T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm), T.fpar(e, T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), fmap(h, dm, List.append(&1, T.FS, l, r)), app2(Array<F32>, T.fpar(e, T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm), fmap(h, dm, l), T.fpar(e, T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm), fmap(h, dm, r), q1, q2), Equal.sym(List<&1, Array<F32>>, fmap(h, dm, List.append(&1, T.FS, l, r)), List.append(&1, Array<F32>, fmap(h, dm, l), fmap(h, dm, r)), bb))}
392
393def m1b(-e: Nat, -l: List<&1, T.FS>, -r: List<&1, T.FS>, +n: Nat, +h: Nat, +dm: T.Dm, -hh1: T.Halves<T.FS>, s1: {hh1 == T.cut(T.FS, TA.half(TA.half(n)), l, Nil{}) : T.Halves<T.FS>}, cat1: {T.hcat(T.FS, hh1) == l : List<&1, T.FS>}, -hh2: T.Halves<T.FS>, s2: {hh2 == T.cut(T.FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}) : T.Halves<T.FS>}, cat2: {T.hcat(T.FS, hh2) == r : List<&1, T.FS>}, r1: MR<e, hh1, TA.half(n), h, dm>, r2: MR<e, hh2, Nat.sub(n, TA.half(n)), h, dm>) -> MR<1n+e, T.Halves{l, r}, n, h, dm>:
394 match r1 r2:
395 case MR{y1, sy1, f1} MR{y2, sy2, f2}:
396 m1c(e, l, r, n, h, dm, hh1, s1, cat1, hh2, s2, cat2, y1, sy1, f1, y2, sy2, f2, map_app(y1, h, dm, y2))
397
398def m1a(-e: Nat, -l: List<&1, T.FS>, -r: List<&1, T.FS>, +n: Nat, +h: Nat, +dm: T.Dm, c1: CGR<TA.half(TA.half(n)), l, Nil{}>, c2: CGR<TA.half(Nat.sub(n, TA.half(n))), r, Nil{}>, k1: @g: T.Halves<T.FS> -> MR<e, g, TA.half(n), h, dm>, k2: @g: T.Halves<T.FS> -> MR<e, g, Nat.sub(n, TA.half(n)), h, dm>) -> MR<1n+e, T.Halves{l, r}, n, h, dm>:
399 match c1 c2:
400 case CGR{hh1, s1, cat1} CGR{hh2, s2, cat2}:
401 m1b(e, l, r, n, h, dm, hh1, s1, cat1, hh2, s2, cat2, k1(hh1), k2(hh2))
402
403def tree(d: Nat, gs: T.Halves<T.FS>, +n: Nat, +h: Nat, +dm: T.Dm) -> MR<d, gs, n, h, dm>:
404 match d gs:
405 case 0n T.Halves{l, r}:
406 m0a(l, r, n, h, dm, l1(l, h, dm, Nil{}))
407 case 1n++e T.Halves{l, r}:
408 m1a(e, l, r, n, h, dm, cg(TA.half(TA.half(n)), l, Nil{}), cg(TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), g1 => tree(e, g1, TA.half(n), h, dm), g2 => tree(e, g2, Nat.sub(n, TA.half(n)), h, dm))
409
410# the law: the tree gives what fleaf, the loop over the sequences one after another, gives
411def par_fin(-d: Nat, -gs: T.Halves<T.FS>, -n: Nat, +h: Nat, +dm: T.Dm, -ys: List<&1, T.FS>, +sy: {ys == T.hcat(T.FS, gs) : List<&1, T.FS>}, f: {T.fpar(d, gs, n, h, dm) == fmap(h, dm, T.hcat(T.FS, gs)) : List<&1, Array<F32>>}, r: L1R<h, dm, ys, Nil{}>) -> {T.fpar(d, gs, n, h, dm) == T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}) : List<&1, Array<F32>>}:
412 match r:
413 case L1R{ys2, s2, fl}:
414 Equal.trans(List<&1, Array<F32>>, T.fpar(d, gs, n, h, dm), fmap(h, dm, T.hcat(T.FS, gs)), T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}), f,
415 Equal.sym(List<&1, Array<F32>>, T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}), fmap(h, dm, T.hcat(T.FS, gs)),
416 Equal.trans(List<&1, Array<F32>>, T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}), T.fleaf(ys, h, dm, Nil{}), fmap(h, dm, T.hcat(T.FS, gs)), Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => T.fleaf(z, h, dm, Nil{}), T.hcat(T.FS, gs), ys, Equal.sym(List<&1, T.FS>, ys, T.hcat(T.FS, gs), sy)),
417 Equal.trans(List<&1, Array<F32>>, T.fleaf(ys, h, dm, Nil{}), fmap(h, dm, ys), fmap(h, dm, T.hcat(T.FS, gs)), fl, Equal.cong(List<&1, T.FS>, List<&1, Array<F32>>, z => fmap(h, dm, z), ys, T.hcat(T.FS, gs), sy)))))
418
419def par_mid(-d: Nat, -gs: T.Halves<T.FS>, -n: Nat, +h: Nat, +dm: T.Dm, r: MR<d, gs, n, h, dm>) -> {T.fpar(d, gs, n, h, dm) == T.fleaf(T.hcat(T.FS, gs), h, dm, Nil{}) : List<&1, Array<F32>>}:
420 match r:
421 case MR{ys, sy, f}:
422 par_fin(d, gs, n, h, dm, ys, sy, f, l1(ys, h, dm, Nil{}))
423
424# PROOF (attn_par_fwd): by induction on the depth, each fact handing its list back (the lists hold arrays, so each is
425# used once): a leaf pair is fleaf twice, and fleaf is the map pushed onto its accumulator (l1); the map of two lists
426# appended is the two maps appended (map_app); cut's halves, appended, are the list it cut (cg); a deeper tree is two
427# trees on cut's halves (the closures).
428def Laws.attn_par_fwd(d, gs, n, h, dm):
429 par_mid(d, gs, n, h, dm, tree(d, gs, n, h, dm))
430
431# ---- the sequential map the b_tree is compared with: f(x) = fs_o(fheads(h, 0, dm, x)) for each x, in order ----
432
433def b_fmap(+h: Nat, +dm: T.Dm, xs: List<&1, T.BG>) -> List<&1, T.DB>:
434 match xs:
435 case Nil{}:
436 Nil{}
437 case Con{x, t}:
438 T.db_of(T.bgroups(h, 0n, dm, x)) <> b_fmap(h, dm, t)
439
440# ---- fleaf (the code's loop with an accumulator) is the map pushed onto the accumulator ----
441
442type b_L1R<-h: Nat, -dm: T.Dm, -xs: List<&1, T.BG>, -acc: List<&1, T.DB>> is Type:
443 b_L1R{ys: List<&1, T.BG>, same: {ys == xs : List<&1, T.BG>}, fact: {T.bleaf(xs, h, dm, acc) == List.reverse.go(&1, T.DB, acc, b_fmap(h, dm, xs)) : List<&1, T.DB>}}
444
445def b_l1_con(x: T.BG, -t: List<&1, T.BG>, +h: Nat, +dm: T.Dm, -acc: List<&1, T.DB>, r: b_L1R<h, dm, t, T.db_of(T.bgroups(h, 0n, dm, x)) <> acc>) -> b_L1R<h, dm, x <> t, acc>:
446 match r:
447 case b_L1R{ys, same, fact}:
448 b_L1R{x <> ys, Equal.cong(List<&1, T.BG>, List<&1, T.BG>, z => x <> z, ys, t, same), fact}
449
450def b_l1(xs: List<&1, T.BG>, +h: Nat, +dm: T.Dm, -acc: List<&1, T.DB>) -> b_L1R<h, dm, xs, acc>:
451 match xs:
452 case Nil{}:
453 b_L1R{Nil{}, {==}, {==}}
454 case Con{x, t}:
455 b_l1_con(x, t, h, dm, acc, b_l1(t, h, dm, T.db_of(T.bgroups(h, 0n, dm, x)) <> acc))
456
457# ---- the map of two lists appended is the two maps appended ----
458
459type b_MAR<-h: Nat, -dm: T.Dm, -a: List<&1, T.BG>, -b: List<&1, T.BG>> is Type:
460 b_MAR{ys: List<&1, T.BG>, same: {ys == a : List<&1, T.BG>}, fact: {b_fmap(h, dm, List.append(&1, T.BG, a, b)) == List.append(&1, T.DB, b_fmap(h, dm, a), b_fmap(h, dm, b)) : List<&1, T.DB>}}
461
462def b_ma_con(x: T.BG, -t: List<&1, T.BG>, +h: Nat, +dm: T.Dm, -b: List<&1, T.BG>, r: b_MAR<h, dm, t, b>) -> b_MAR<h, dm, x <> t, b>:
463 match r:
464 case b_MAR{ys, same, fact}:
465 b_MAR{x <> ys, Equal.cong(List<&1, T.BG>, List<&1, T.BG>, z => x <> z, ys, t, same), Equal.cong(List<&1, T.DB>, List<&1, T.DB>, z => T.db_of(T.bgroups(h, 0n, dm, x)) <> z, b_fmap(h, dm, List.append(&1, T.BG, t, b)), List.append(&1, T.DB, b_fmap(h, dm, t), b_fmap(h, dm, b)), fact)}
466
467def b_map_app(a: List<&1, T.BG>, +h: Nat, +dm: T.Dm, -b: List<&1, T.BG>) -> b_MAR<h, dm, a, b>:
468 match a:
469 case Nil{}:
470 b_MAR{Nil{}, {==}, {==}}
471 case Con{x, t}:
472 b_ma_con(x, t, h, dm, b, b_map_app(t, h, dm, b))
473
474# ---- a reversed onto b, then c appended: a reversed onto b ++ c ----
475
476type b_GAR<-a: List<&1, T.BG>, -b: List<&1, T.BG>, -c: List<&1, T.BG>> is Type:
477 b_GAR{ys: List<&1, T.BG>, same: {ys == a : List<&1, T.BG>}, fact: {List.append(&1, T.BG, List.reverse.go(&1, T.BG, a, b), c) == List.reverse.go(&1, T.BG, a, List.append(&1, T.BG, b, c)) : List<&1, T.BG>}}
478
479def b_ga_con(x: T.BG, -t: List<&1, T.BG>, -b: List<&1, T.BG>, -c: List<&1, T.BG>, r: b_GAR<t, x <> b, c>) -> b_GAR<x <> t, b, c>:
480 match r:
481 case b_GAR{ys, same, fact}:
482 b_GAR{x <> ys, Equal.cong(List<&1, T.BG>, List<&1, T.BG>, z => x <> z, ys, t, same), fact}
483
484def b_go_app(a: List<&1, T.BG>, -b: List<&1, T.BG>, -c: List<&1, T.BG>) -> b_GAR<a, b, c>:
485 match a:
486 case Nil{}:
487 b_GAR{Nil{}, {==}, {==}}
488 case Con{x, t}:
489 b_ga_con(x, t, b, c, b_go_app(t, x <> b, c))
490
491# ---- cut gives two halves that, appended, are acc reversed onto xs ----
492
493type b_CGR<-k: Nat, -xs: List<&1, T.BG>, -acc: List<&1, T.BG>> is Type:
494 b_CGR{hh: T.Halves<T.BG>, same: {hh == T.cut(T.BG, k, xs, acc) : T.Halves<T.BG>}, cat: {T.hcat(T.BG, hh) == List.reverse.go(&1, T.BG, acc, xs) : List<&1, T.BG>}}
495
496def b_cg_base(-k: Nat, xs: List<&1, T.BG>, -acc: List<&1, T.BG>, e: {T.cut(T.BG, k, xs, acc) == T.Halves{List.reverse(&1, T.BG, acc), xs} : T.Halves<T.BG>}, r: b_GAR<acc, Nil{}, xs>) -> b_CGR<k, xs, acc>:
497 match r:
498 case b_GAR{ys, +same, fact}:
499 b_CGR{T.Halves{List.reverse(&1, T.BG, ys), xs},
500 Equal.trans(T.Halves<T.BG>, T.Halves{List.reverse(&1, T.BG, ys), xs}, T.Halves{List.reverse(&1, T.BG, acc), xs}, T.cut(T.BG, k, xs, acc), Equal.cong(List<&1, T.BG>, T.Halves<T.BG>, z => T.Halves{List.reverse(&1, T.BG, z), xs}, ys, acc, same), Equal.sym(T.Halves<T.BG>, T.cut(T.BG, k, xs, acc), T.Halves{List.reverse(&1, T.BG, acc), xs}, e)),
501 Equal.trans(List<&1, T.BG>, List.append(&1, T.BG, List.reverse.go(&1, T.BG, ys, Nil{}), xs), List.append(&1, T.BG, List.reverse.go(&1, T.BG, acc, Nil{}), xs), List.reverse.go(&1, T.BG, acc, xs), Equal.cong(List<&1, T.BG>, List<&1, T.BG>, z => List.append(&1, T.BG, List.reverse.go(&1, T.BG, z, Nil{}), xs), ys, acc, same), fact)}
502
503def b_cg_con(-m: Nat, -x: T.BG, -t: List<&1, T.BG>, -acc: List<&1, T.BG>, r: b_CGR<m, t, x <> acc>) -> b_CGR<1n+m, x <> t, acc>:
504 match r:
505 case b_CGR{hh, same, cat}:
506 b_CGR{hh, same, cat}
507
508def b_cg(k: Nat, xs: List<&1, T.BG>, acc: List<&1, T.BG>) -> b_CGR<k, xs, acc>:
509 match k xs:
510 case 1n+m Con{x, t}:
511 b_cg_con(m, x, t, acc, b_cg(m, t, x <> acc))
512 case 0n ys:
513 b_cg_base(0n, ys, acc, {==}, b_go_app(acc, Nil{}, ys))
514 case 1n+m Nil{}:
515 b_cg_base(1n+m, Nil{}, acc, {==}, b_go_app(acc, Nil{}, Nil{}))
516
517# ---- the b_tree: fpar at depth d on two halves is the map of the halves appended ----
518
519type b_MR<-d: Nat, -gs: T.Halves<T.BG>, -n: Nat, -h: Nat, -dm: T.Dm> is Type:
520 b_MR{ys: List<&1, T.BG>, same: {ys == T.hcat(T.BG, gs) : List<&1, T.BG>}, fact: {T.bpar(d, gs, n, h, dm) == b_fmap(h, dm, T.hcat(T.BG, gs)) : List<&1, T.DB>}}
521
522# depth 0: the two leaves' loops, appended
523def b_m0c(-l: List<&1, T.BG>, -r: List<&1, T.BG>, -n: Nat, +h: Nat, +dm: T.Dm, -l1y: List<&1, T.BG>, +sl: {l1y == l : List<&1, T.BG>}, fl: {T.bleaf(l, h, dm, Nil{}) == b_fmap(h, dm, l) : List<&1, T.DB>}, r1: List<&1, T.BG>, sr: {r1 == r : List<&1, T.BG>}, fr: {T.bleaf(r, h, dm, Nil{}) == b_fmap(h, dm, r) : List<&1, T.DB>}, c: b_MAR<h, dm, l1y, r>) -> b_MR<0n, T.Halves{l, r}, n, h, dm>:
524 match c:
525 case b_MAR{l2, s2, fa}:
526 b_MR{List.append(&1, T.BG, l2, r1), app2(T.BG, l2, l, r1, r, Equal.trans(List<&1, T.BG>, l2, l1y, l, s2, sl), sr),
527 Equal.trans(List<&1, T.DB>, List.append(&1, T.DB, T.bleaf(l, h, dm, Nil{}), T.bleaf(r, h, dm, Nil{})), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), b_fmap(h, dm, List.append(&1, T.BG, l, r)), app2(T.DB, T.bleaf(l, h, dm, Nil{}), b_fmap(h, dm, l), T.bleaf(r, h, dm, Nil{}), b_fmap(h, dm, r), fl, fr),
528 Equal.sym(List<&1, T.DB>, b_fmap(h, dm, List.append(&1, T.BG, l, r)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)),
529 Equal.trans(List<&1, T.DB>, b_fmap(h, dm, List.append(&1, T.BG, l, r)), b_fmap(h, dm, List.append(&1, T.BG, l1y, r)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, List.append(&1, T.BG, z, r)), l, l1y, Equal.sym(List<&1, T.BG>, l1y, l, sl)),
530 Equal.trans(List<&1, T.DB>, b_fmap(h, dm, List.append(&1, T.BG, l1y, r)), List.append(&1, T.DB, b_fmap(h, dm, l1y), b_fmap(h, dm, r)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), fa, Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => List.append(&1, T.DB, b_fmap(h, dm, z), b_fmap(h, dm, r)), l1y, l, sl)))))}
531
532def b_m0b(-l: List<&1, T.BG>, -r: List<&1, T.BG>, -n: Nat, +h: Nat, +dm: T.Dm, l1y: List<&1, T.BG>, +sl: {l1y == l : List<&1, T.BG>}, fl: {T.bleaf(l, h, dm, Nil{}) == b_fmap(h, dm, l) : List<&1, T.DB>}, b: b_L1R<h, dm, r, Nil{}>) -> b_MR<0n, T.Halves{l, r}, n, h, dm>:
533 match b:
534 case b_L1R{r1, sr, fr}:
535 b_m0c(l, r, n, h, dm, l1y, sl, fl, r1, sr, fr, b_map_app(l1y, h, dm, r))
536
537def b_m0a(-l: List<&1, T.BG>, r: List<&1, T.BG>, -n: Nat, +h: Nat, +dm: T.Dm, a: b_L1R<h, dm, l, Nil{}>) -> b_MR<0n, T.Halves{l, r}, n, h, dm>:
538 match a:
539 case b_L1R{l1y, sl, fl}:
540 b_m0b(l, r, n, h, dm, l1y, sl, fl, b_l1(r, h, dm, Nil{}))
541
542# depth 1 + e: each half cut in two and its b_tree (the closures run main at depth e), then the two maps appended
543def b_m1c(-e: Nat, -l: List<&1, T.BG>, -r: List<&1, T.BG>, +n: Nat, +h: Nat, +dm: T.Dm, -hh1: T.Halves<T.BG>, s1: {hh1 == T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}) : T.Halves<T.BG>}, +cat1: {T.hcat(T.BG, hh1) == l : List<&1, T.BG>}, -hh2: T.Halves<T.BG>, s2: {hh2 == T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}) : T.Halves<T.BG>}, +cat2: {T.hcat(T.BG, hh2) == r : List<&1, T.BG>}, -y1: List<&1, T.BG>, sy1: {y1 == T.hcat(T.BG, hh1) : List<&1, T.BG>}, f1: {T.bpar(e, hh1, TA.half(n), h, dm) == b_fmap(h, dm, T.hcat(T.BG, hh1)) : List<&1, T.DB>}, y2: List<&1, T.BG>, sy2: {y2 == T.hcat(T.BG, hh2) : List<&1, T.BG>}, f2: {T.bpar(e, hh2, Nat.sub(n, TA.half(n)), h, dm) == b_fmap(h, dm, T.hcat(T.BG, hh2)) : List<&1, T.DB>}, ma: b_MAR<h, dm, y1, y2>) -> b_MR<1n+e, T.Halves{l, r}, n, h, dm>:
544 match ma:
545 case b_MAR{y1b, sb, fa}:
546 +ly1 = Equal.trans(List<&1, T.BG>, y1, T.hcat(T.BG, hh1), l, sy1, cat1)
547 +ly2 = Equal.trans(List<&1, T.BG>, y2, T.hcat(T.BG, hh2), r, sy2, cat2)
548 +q1 = Equal.trans(List<&1, T.DB>, T.bpar(e, T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm), T.bpar(e, hh1, TA.half(n), h, dm), b_fmap(h, dm, l), Equal.cong(T.Halves<T.BG>, List<&1, T.DB>, z => T.bpar(e, z, TA.half(n), h, dm), T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}), hh1, Equal.sym(T.Halves<T.BG>, hh1, T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}), s1)), Equal.trans(List<&1, T.DB>, T.bpar(e, hh1, TA.half(n), h, dm), b_fmap(h, dm, T.hcat(T.BG, hh1)), b_fmap(h, dm, l), f1, Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, z), T.hcat(T.BG, hh1), l, cat1)))
549 +q2 = Equal.trans(List<&1, T.DB>, T.bpar(e, T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm), T.bpar(e, hh2, Nat.sub(n, TA.half(n)), h, dm), b_fmap(h, dm, r), Equal.cong(T.Halves<T.BG>, List<&1, T.DB>, z => T.bpar(e, z, Nat.sub(n, TA.half(n)), h, dm), T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), hh2, Equal.sym(T.Halves<T.BG>, hh2, T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), s2)), Equal.trans(List<&1, T.DB>, T.bpar(e, hh2, Nat.sub(n, TA.half(n)), h, dm), b_fmap(h, dm, T.hcat(T.BG, hh2)), b_fmap(h, dm, r), f2, Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, z), T.hcat(T.BG, hh2), r, cat2)))
550 +bb = Equal.trans(List<&1, T.DB>, b_fmap(h, dm, List.append(&1, T.BG, l, r)), b_fmap(h, dm, List.append(&1, T.BG, y1, y2)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, z), List.append(&1, T.BG, l, r), List.append(&1, T.BG, y1, y2), app2(T.BG, l, y1, r, y2, Equal.sym(List<&1, T.BG>, y1, l, ly1), Equal.sym(List<&1, T.BG>, y2, r, ly2))), Equal.trans(List<&1, T.DB>, b_fmap(h, dm, List.append(&1, T.BG, y1, y2)), List.append(&1, T.DB, b_fmap(h, dm, y1), b_fmap(h, dm, y2)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), fa, app2(T.DB, b_fmap(h, dm, y1), b_fmap(h, dm, l), b_fmap(h, dm, y2), b_fmap(h, dm, r), Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, z), y1, l, ly1), Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, z), y2, r, ly2))))
551 b_MR{List.append(&1, T.BG, y1b, y2), app2(T.BG, y1b, l, y2, r, Equal.trans(List<&1, T.BG>, y1b, y1, l, sb, ly1), ly2),
552 Equal.trans(List<&1, T.DB>, List.append(&1, T.DB, T.bpar(e, T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm), T.bpar(e, T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), b_fmap(h, dm, List.append(&1, T.BG, l, r)), app2(T.DB, T.bpar(e, T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm), b_fmap(h, dm, l), T.bpar(e, T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm), b_fmap(h, dm, r), q1, q2), Equal.sym(List<&1, T.DB>, b_fmap(h, dm, List.append(&1, T.BG, l, r)), List.append(&1, T.DB, b_fmap(h, dm, l), b_fmap(h, dm, r)), bb))}
553
554def b_m1b(-e: Nat, -l: List<&1, T.BG>, -r: List<&1, T.BG>, +n: Nat, +h: Nat, +dm: T.Dm, -hh1: T.Halves<T.BG>, s1: {hh1 == T.cut(T.BG, TA.half(TA.half(n)), l, Nil{}) : T.Halves<T.BG>}, cat1: {T.hcat(T.BG, hh1) == l : List<&1, T.BG>}, -hh2: T.Halves<T.BG>, s2: {hh2 == T.cut(T.BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}) : T.Halves<T.BG>}, cat2: {T.hcat(T.BG, hh2) == r : List<&1, T.BG>}, r1: b_MR<e, hh1, TA.half(n), h, dm>, r2: b_MR<e, hh2, Nat.sub(n, TA.half(n)), h, dm>) -> b_MR<1n+e, T.Halves{l, r}, n, h, dm>:
555 match r1 r2:
556 case b_MR{y1, sy1, f1} b_MR{y2, sy2, f2}:
557 b_m1c(e, l, r, n, h, dm, hh1, s1, cat1, hh2, s2, cat2, y1, sy1, f1, y2, sy2, f2, b_map_app(y1, h, dm, y2))
558
559def b_m1a(-e: Nat, -l: List<&1, T.BG>, -r: List<&1, T.BG>, +n: Nat, +h: Nat, +dm: T.Dm, c1: b_CGR<TA.half(TA.half(n)), l, Nil{}>, c2: b_CGR<TA.half(Nat.sub(n, TA.half(n))), r, Nil{}>, k1: @g: T.Halves<T.BG> -> b_MR<e, g, TA.half(n), h, dm>, k2: @g: T.Halves<T.BG> -> b_MR<e, g, Nat.sub(n, TA.half(n)), h, dm>) -> b_MR<1n+e, T.Halves{l, r}, n, h, dm>:
560 match c1 c2:
561 case b_CGR{hh1, s1, cat1} b_CGR{hh2, s2, cat2}:
562 b_m1b(e, l, r, n, h, dm, hh1, s1, cat1, hh2, s2, cat2, k1(hh1), k2(hh2))
563
564def b_tree(d: Nat, gs: T.Halves<T.BG>, +n: Nat, +h: Nat, +dm: T.Dm) -> b_MR<d, gs, n, h, dm>:
565 match d gs:
566 case 0n T.Halves{l, r}:
567 b_m0a(l, r, n, h, dm, b_l1(l, h, dm, Nil{}))
568 case 1n++e T.Halves{l, r}:
569 b_m1a(e, l, r, n, h, dm, b_cg(TA.half(TA.half(n)), l, Nil{}), b_cg(TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), g1 => b_tree(e, g1, TA.half(n), h, dm), g2 => b_tree(e, g2, Nat.sub(n, TA.half(n)), h, dm))
570
571# the law: the b_tree gives what fleaf, the loop over the sequences one after another, gives
572def b_par_fin(-d: Nat, -gs: T.Halves<T.BG>, -n: Nat, +h: Nat, +dm: T.Dm, -ys: List<&1, T.BG>, +sy: {ys == T.hcat(T.BG, gs) : List<&1, T.BG>}, f: {T.bpar(d, gs, n, h, dm) == b_fmap(h, dm, T.hcat(T.BG, gs)) : List<&1, T.DB>}, r: b_L1R<h, dm, ys, Nil{}>) -> {T.bpar(d, gs, n, h, dm) == T.bleaf(T.hcat(T.BG, gs), h, dm, Nil{}) : List<&1, T.DB>}:
573 match r:
574 case b_L1R{ys2, s2, fl}:
575 Equal.trans(List<&1, T.DB>, T.bpar(d, gs, n, h, dm), b_fmap(h, dm, T.hcat(T.BG, gs)), T.bleaf(T.hcat(T.BG, gs), h, dm, Nil{}), f,
576 Equal.sym(List<&1, T.DB>, T.bleaf(T.hcat(T.BG, gs), h, dm, Nil{}), b_fmap(h, dm, T.hcat(T.BG, gs)),
577 Equal.trans(List<&1, T.DB>, T.bleaf(T.hcat(T.BG, gs), h, dm, Nil{}), T.bleaf(ys, h, dm, Nil{}), b_fmap(h, dm, T.hcat(T.BG, gs)), Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => T.bleaf(z, h, dm, Nil{}), T.hcat(T.BG, gs), ys, Equal.sym(List<&1, T.BG>, ys, T.hcat(T.BG, gs), sy)),
578 Equal.trans(List<&1, T.DB>, T.bleaf(ys, h, dm, Nil{}), b_fmap(h, dm, ys), b_fmap(h, dm, T.hcat(T.BG, gs)), fl, Equal.cong(List<&1, T.BG>, List<&1, T.DB>, z => b_fmap(h, dm, z), ys, T.hcat(T.BG, gs), sy)))))
579
580def b_par_mid(-d: Nat, -gs: T.Halves<T.BG>, -n: Nat, +h: Nat, +dm: T.Dm, r: b_MR<d, gs, n, h, dm>) -> {T.bpar(d, gs, n, h, dm) == T.bleaf(T.hcat(T.BG, gs), h, dm, Nil{}) : List<&1, T.DB>}:
581 match r:
582 case b_MR{ys, sy, f}:
583 b_par_fin(d, gs, n, h, dm, ys, sy, f, b_l1(ys, h, dm, Nil{}))
584
585# PROOF (attn_par_bwd): attn_par_fwd's proof, for BG and bleaf.
586def Laws.attn_par_bwd(d, gs, n, h, dm):
587 b_par_mid(d, gs, n, h, dm, b_tree(d, gs, n, h, dm))
588
589# =====================================================================================================
590# Grouped-query heads (llama.bend's proof)
591# =====================================================================================================
592
593# one block of 1 + m numbers adds 1 to the quotient (Nat.divmod.go counts down the remainder from m)
594def div_block(+m: Nat, -n: Nat, -d: Nat, -r: Nat) -> {Nat.divmod.go(Nat.add(1n+m, n), m, d, r) == Nat.divmod.go(n, Nat.add(m, r), 1n+d, 0n) : Nat & Nat}:
595 match m:
596 case 0n:
597 {==}
598 case 1n+p:
599 %Equal.sym(Nat, 1n+Nat.add(p, r), Nat.add(p, 1n+r), NL.add_succ(p, r)) : {Nat.divmod.go(Nat.add(1n+p, n), p, d, 1n+r) == Nat.divmod.go(n, _, 1n+d, 0n) : Nat & Nat}
600 div_block(p, n, d, 1n+r)
601
602# j whole blocks of 1 + s numbers add j to the quotient
603def div_blocks(+j: Nat, +s: Nat, +n: Nat, +d: Nat) -> {Nat.divmod.go(Nat.add(Nat.mul(j, 1n+s), n), s, d, 0n) == Nat.divmod.go(n, s, Nat.add(j, d), 0n) : Nat & Nat}:
604 match j:
605 case 0n:
606 {==}
607 case 1n+jp:
608 %NL.add_assoc(1n+s, Nat.mul(jp, 1n+s), n) : {Nat.divmod.go(_, s, d, 0n) == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
609 %Equal.sym(Nat & Nat, Nat.divmod.go(Nat.add(1n+s, Nat.add(Nat.mul(jp, 1n+s), n)), s, d, 0n), Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), Nat.add(s, 0n), 1n+d, 0n), div_block(s, Nat.add(Nat.mul(jp, 1n+s), n), d, 0n)) : {_ == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
610 %Equal.sym(Nat, Nat.add(s, 0n), s, NL.add_zero(s)) : {Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), _, 1n+d, 0n) == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
611 %Equal.sym(Nat, 1n+Nat.add(jp, d), Nat.add(jp, 1n+d), NL.add_succ(jp, d)) : {Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), s, 1n+d, 0n) == Nat.divmod.go(n, s, _, 0n) : Nat & Nat}
612 div_blocks(jp, s, n, 1n+d)
613
614# fewer than 1 + s numbers after the last whole block leave the quotient alone
615def div_rest(+i: Nat, -k: Nat, -d: Nat, -r: Nat) -> {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), d, r)) == d : Nat}:
616 match i:
617 case 0n:
618 {==}
619 case 1n+p:
620 div_rest(p, k, d, 1n+r)
621
622# PROOF (gqa_head): Nat.div counts how many whole blocks of g numbers fit. j·g is j whole blocks (div_blocks), and
623# i < g numbers after them do not complete another one (div_rest).
624def Laws.gqa_head(j, i, k):
625 %Equal.sym(Nat & Nat, Nat.divmod.go(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), Nat.add(i, k), 0n, 0n), Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n), div_blocks(j, Nat.add(i, k), i, 0n)) : {Pair.fst(Nat, Nat, _) == j : Nat}
626 %NL.add_zero(j) : {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n)) == _ : Nat}
627 div_rest(i, k, Nat.add(j, 0n), 0n)
628
629# PROOF (kv_of_group): kv_of is Nat.div, so this is gqa_head.
630def Laws.kv_of_group(j, i, k):
631 Laws.gqa_head(j, i, k)
632
633# PROOF (kv_cols): head_cols' key column is kv_of(head, g)·hd; kv_of of a head of group j is j (kv_of_group).
634def Laws.kv_cols(t, hd, w, kvw, sc, j, i, k):
635 Equal.cong(Nat, U32, n => U32.from_nat(Nat.mul(n, hd)), T.kv_of(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)), j, Laws.kv_of_group(j, i, k))
636
637# =====================================================================================================
638# Tensors: shapes
639# =====================================================================================================
640
641# a || b == b || a, and a && b == b && a: every case computes
642def or_comm(+a: Bool, +b: Bool) -> {Bool.or(a, b) == Bool.or(b, a) : Bool}:
643 match a b:
644 case True{} True{}:
645 {==}
646 case True{} False{}:
647 {==}
648 case False{} True{}:
649 {==}
650 case False{} False{}:
651 {==}
652
653def and_comm(+a: Bool, +b: Bool) -> {Bool.and(a, b) == Bool.and(b, a) : Bool}:
654 match a b:
655 case True{} True{}:
656 {==}
657 case True{} False{}:
658 {==}
659 case False{} True{}:
660 {==}
661 case False{} False{}:
662 {==}
663
664# equal operands give equal conjunctions and disjunctions
665def and_cong(+p: Bool, +p2: Bool, +q: Bool, +q2: Bool, e: {p == p2 : Bool}, f: {q == q2 : Bool}) -> {Bool.and(p, q) == Bool.and(p2, q2) : Bool}:
666 Equal.trans(Bool, Bool.and(p, q), Bool.and(p2, q), Bool.and(p2, q2), Equal.cong(Bool, Bool, x => Bool.and(x, q), p, p2, e), Equal.cong(Bool, Bool, x => Bool.and(p2, x), q, q2, f))
667
668def or_cong(+p: Bool, +p2: Bool, +q: Bool, +q2: Bool, e: {p == p2 : Bool}, f: {q == q2 : Bool}) -> {Bool.or(p, q) == Bool.or(p2, q2) : Bool}:
669 Equal.trans(Bool, Bool.or(p, q), Bool.or(p2, q), Bool.or(p2, q2), Equal.cong(Bool, Bool, x => Bool.or(x, q), p, p2, e), Equal.cong(Bool, Bool, x => Bool.or(p2, x), q, q2, f))
670
671# one axis: bc_dim is written as pair(x, y) || pair(y, x), so swapping x and y only swaps the two sides of the ||
672def bc_dim_sym(+x: Nat, +y: Nat, +z: Nat) -> {TT.bc_dim(x, y, z) == TT.bc_dim(y, x, z) : Bool}:
673 or_comm(TT.bc_pair(x, y, z), TT.bc_pair(y, x, z))
674
675# axis by axis: the same on each axis (bc_dim_sym), and by induction on the rest; shapes of different ranks are False
676# either way round
677def bc_eq_sym(a: List<&2, Nat>, b: List<&2, Nat>, o: List<&2, Nat>) -> {TT.bc_eq(a, b, o) == TT.bc_eq(b, a, o) : Bool}:
678 match a b o:
679 case Nil{} Nil{} Nil{}:
680 {==}
681 case Nil{} Nil{} Con{z, ro}:
682 {==}
683 case Nil{} Con{y, rb} Nil{}:
684 {==}
685 case Nil{} Con{y, rb} Con{z, ro}:
686 {==}
687 case Con{x, ra} Nil{} Nil{}:
688 {==}
689 case Con{x, ra} Nil{} Con{z, ro}:
690 {==}
691 case Con{x, ra} Con{y, rb} Nil{}:
692 {==}
693 case Con{+x, +ra} Con{+y, +rb} Con{+z, +ro}:
694 and_cong(TT.bc_dim(x, y, z), TT.bc_dim(y, x, z), TT.bc_eq(ra, rb, ro), TT.bc_eq(rb, ra, ro), bc_dim_sym(x, y, z), bc_eq_sym(ra, rb, ro))
695
696# the ranks: the same two comparisons and the same two equalities, in the other order
697def bc_rank_sym(+la: Nat, +lb: Nat, +lo: Nat) -> {TT.bc_rank(la, lb, lo) == TT.bc_rank(lb, la, lo) : Bool}:
698 and_cong(Bool.and(Nat.is_le(la, lo), Nat.is_le(lb, lo)), Bool.and(Nat.is_le(lb, lo), Nat.is_le(la, lo)), Bool.or(Nat.is_eq(la, lo), Nat.is_eq(lb, lo)), Bool.or(Nat.is_eq(lb, lo), Nat.is_eq(la, lo)),
699 and_comm(Nat.is_le(la, lo), Nat.is_le(lb, lo)), or_comm(Nat.is_eq(la, lo), Nat.is_eq(lb, lo)))
700
701def Laws.bc_sym(a, b, o):
702 and_cong(TT.bc_rank(List.length(&2, Nat, a), List.length(&2, Nat, b), List.length(&2, Nat, o)), TT.bc_rank(List.length(&2, Nat, b), List.length(&2, Nat, a), List.length(&2, Nat, o)),
703 TT.bc_eq(TT.pad(a, List.length(&2, Nat, o)), TT.pad(b, List.length(&2, Nat, o)), o), TT.bc_eq(TT.pad(b, List.length(&2, Nat, o)), TT.pad(a, List.length(&2, Nat, o)), o),
704 bc_rank_sym(List.length(&2, Nat, a), List.length(&2, Nat, b), List.length(&2, Nat, o)), bc_eq_sym(TT.pad(a, List.length(&2, Nat, o)), TT.pad(b, List.length(&2, Nat, o)), o))
705
706# h * (a * r) == a * (h * r): associativity, then commutativity of the first two factors, then associativity back
707def mul_swap(+h: Nat, +a: Nat, +r: Nat) -> {Nat.mul(h, Nat.mul(a, r)) == Nat.mul(a, Nat.mul(h, r)) : Nat}:
708 Equal.trans(Nat, Nat.mul(h, Nat.mul(a, r)), Nat.mul(Nat.mul(h, a), r), Nat.mul(a, Nat.mul(h, r)), NL.mul_assoc(h, a, r),
709 Equal.trans(Nat, Nat.mul(Nat.mul(h, a), r), Nat.mul(Nat.mul(a, h), r), Nat.mul(a, Nat.mul(h, r)),
710 Equal.cong(Nat, Nat, x => Nat.mul(x, r), Nat.mul(h, a), Nat.mul(a, h), NL.mul_comm(h, a)),
711 Equal.sym(Nat, Nat.mul(a, Nat.mul(h, r)), Nat.mul(Nat.mul(a, h), r), NL.mul_assoc(a, h, r))))
712
713# by induction on the shape and the axis: axis 0 computes; past it, the first size h stays in front on both sides, and
714# mul_swap moves it past size(k)
715def Laws.sum_numel(s, k):
716 match s k:
717 case Nil{} _:
718 {==}
719 case Con{h, t} 0n:
720 {==}
721 case Con{+h, +t} 1n+p:
722 Equal.trans(Nat, Nat.mul(h, NL.product(t)), Nat.mul(h, Nat.mul(TT.at1(t, p), NL.product(TT.remove(t, p)))), Nat.mul(TT.at1(t, p), Nat.mul(h, NL.product(TT.remove(t, p)))),
723 Equal.cong(Nat, Nat, x => Nat.mul(h, x), NL.product(t), Nat.mul(TT.at1(t, p), NL.product(TT.remove(t, p))), Laws.sum_numel(t, p)),
724 mul_swap(h, TT.at1(t, p), NL.product(TT.remove(t, p))))
725
726# =====================================================================================================
727# Tensors: narrow, cat, of_mat, permute
728# =====================================================================================================
729
730# by induction: axis 0 is replaced and then taken out; past it, the first size stays in front on both sides
731def Laws.narrow_rest(s, d, l):
732 match s d:
733 case Nil{} _:
734 {==}
735 case Con{h, t} 0n:
736 {==}
737 case Con{+h, +t} 1n+p:
738 Equal.cong(List<&2, Nat>, List<&2, Nat>, x => h <> x, TT.remove(TT.setd(t, p, l), p), TT.remove(t, p), Laws.narrow_rest(t, p, l))
739
740# a shape with no axis d contradicts the premise (d < rank is False); axis 0 is l; past it, the induction
741def Laws.narrow_axis(s, d, l, e):
742 match s d:
743 case Nil{} 0n:
744 Empty.absurd({TT.at1(TT.setd([], 0n, l), 0n) == l : Nat}, false_true(Equal.sym(Bool, True{}, False{}, e)))
745 case Nil{} 1n+p:
746 Empty.absurd({TT.at1(TT.setd([], 1n+p, l), 1n+p) == l : Nat}, false_true(Equal.sym(Bool, True{}, False{}, e)))
747 case Con{h, t} 0n:
748 {==}
749 case Con{h, +t} 1n+p:
750 Laws.narrow_axis(t, p, l, e)
751
752# axis 0: (x + y)·n == x·n + y·n (distributivity); past it, the induction and k·(a + b) == k·a + k·b
753def Laws.cat_numel(s, d, x, y, e):
754 match s d:
755 case Nil{} 0n:
756 Empty.absurd({TT.numel(TT.setd([], 0n, Nat.add(x, y))) == Nat.add(TT.numel(TT.setd([], 0n, x)), TT.numel(TT.setd([], 0n, y))) : Nat}, false_true(Equal.sym(Bool, True{}, False{}, e)))
757 case Nil{} 1n+p:
758 Empty.absurd({TT.numel(TT.setd([], 1n+p, Nat.add(x, y))) == Nat.add(TT.numel(TT.setd([], 1n+p, x)), TT.numel(TT.setd([], 1n+p, y))) : Nat}, false_true(Equal.sym(Bool, True{}, False{}, e)))
759 case Con{h, +t} 0n:
760 Equal.sym(Nat, Nat.add(Nat.mul(x, NL.product(t)), Nat.mul(y, NL.product(t))), Nat.mul(Nat.add(x, y), NL.product(t)), NL.mul_dist(x, y, NL.product(t)))
761 case Con{+h, +t} 1n+p:
762 Equal.trans(Nat, Nat.mul(h, NL.product(TT.setd(t, p, Nat.add(x, y)))), Nat.mul(h, Nat.add(NL.product(TT.setd(t, p, x)), NL.product(TT.setd(t, p, y)))), Nat.add(Nat.mul(h, NL.product(TT.setd(t, p, x))), Nat.mul(h, NL.product(TT.setd(t, p, y)))),
763 Equal.cong(Nat, Nat, v => Nat.mul(h, v), NL.product(TT.setd(t, p, Nat.add(x, y))), Nat.add(NL.product(TT.setd(t, p, x)), NL.product(TT.setd(t, p, y))), Laws.cat_numel(t, p, x, y, e)),
764 mul_dist_l(h, NL.product(TT.setd(t, p, x)), NL.product(TT.setd(t, p, y))))
765
766# r·(c·1) == r·c: c·1 == c
767def Laws.mat_numel(r, c):
768 Equal.cong(Nat, Nat, v => Nat.mul(r, v), Nat.mul(c, 1n), c, NL.mul_one_r(c))
769
770# one axis of the result per entry of p
771def Laws.permute_rank(s, p):
772 match p:
773 case Nil{}:
774 {==}
775 case Con{k, r}:
776 Equal.cong(Nat, Nat, v => 1n+v, TT.rank(TT.permuted(s, r)), List.length(&2, Nat, r), Laws.permute_rank(s, r))
777
778def main() -> Unit:
779 Unit{}