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