main.bend
bendtorch@1.1.0.0 · 3,353 lines · 0 definitions
1# BendTorch: layers with typed forward and backward passes, composed generically, trained in Bend 2.
2#
3# import ../torch/main.bend as T
4#
5# A layer is a record of four functions over its parameters P, its input X, its output Y and what its forward pass
6# saves for the backward pass, S:
7#
8# fwd : P -> X -> Fwd<P, Y, S> the output, the parameters handed back, the saved values
9# bwd : P -> S -> Y -> P -> Bwd<P, X> from dY (typed Y) to dX (typed X), adding dP into a gradient of type P
10# step : Opt -> P -> P -> OptState<P> -> Step<P> one optimizer step on the parameters; the gradient comes back zeroed
11# zeros : Unit -> P a zero P (gradients and optimizer state start from it)
12#
13# The gradient of the input has the input's type and the gradient of the parameters has the parameters' type, so a
14# backward pass that returns a transposed or a wrongly sized gradient does not compile. Layers are passed as templates
15# (`~l`), so `Seq` composes any two layers whose types meet, with no closure called twice and no runtime cost.
16#
17# Matrices are bend-ml's `Mat<r, c>` (bend-ml-tensor-array 0.1.14.0): r*c F32 numbers in one flat Array, row by row,
18# the dimensions erased type parameters. An Array has one owner, so every operation hands back what it reads.
19#
20# What is proved and what is tested: the shapes are in the types (checked); the laws at the end of this file are
21# checked by Bend's kernel; every F32 result is tested against PyTorch by the scripts in reference/, never proved.
22
23import Base
24import bend-ml-tensor-array@0.1.14.0/main.bend as TA
25import bend-ml-nat-lemmas@0.1.2.0/main.bend as NL
26
27# =====================================================================================================
28# The layer protocol
29# =====================================================================================================
30
31# what a forward pass returns: the parameters, the output, what the backward pass will need
32type Fwd<-P: Type, -Y: Type, -S: Type> is Type:
33 Fwd{p: P, y: Y, s: S}
34
35# what a backward pass returns: the parameters, the gradient of the input, the accumulated gradient of the parameters
36type Bwd<-P: Type, -X: Type> is Type:
37 Bwd{p: P, dx: X, g: P}
38
39# a pair of owned values (parameters of two layers or their saved values)
40type Two<-A: Type, -B: Type> is Type:
41 Two{a: A, b: B}
42
43# Type helpers for the left-associated pairs used by Seq/Fork and for Residual's two parameter-free layers.
44def Pair3(A: Type, B: Type, C: Type) -> Type:
45 Two<Two<A, B>, C>
46
47def ResidualType(A: Type) -> Type:
48 Pair3(Unit, A, Unit)
49
50# AdamW's moments stored as packed BF16 (CAdamWBF16): a tree that follows the parameters' pairs as the step walks them,
51# with the buffers of each matrix at its leaf. A packed buffer has half the words of its matrix (two BF16 values a
52# U32, element 2k in the low half of word k, torch/dtype.bend's layout), so it cannot be a value of the parameters'
53# type P: the tree holds the words itself. A leaf is allocated on the matrix's first step (PNew marks a part not yet
54# stepped, and stays at the parameter-free layers), so the full-size buffers never exist. v16 picks v packed as well;
55# otherwise v is F32 (PLeafV). A leaf keeps its matrix's count n, for checkpoints.
56type PTree is Type:
57 PNew{v16: Bool}
58 PLeaf{n: U32, m: Array<U32>, v: Array<U32>}
59 PLeafV{n: U32, m: Array<U32>, v: Array<F32>}
60 PNode{a: PTree, b: PTree}
61
62# Owned optimizer buffers: SGD has none, momentum one, Adam/AdamW two (Packed: two, in packed BF16, see PTree).
63type OptState<-P: Type> is Type:
64 NoMom{}
65 OneMom{m: P}
66 TwoMom{m: P, v: P}
67 Packed{t: PTree}
68
69# Two affine maps are supplied separately: each is called at most once. A packed tree is the same under any relabelling
70# of the parameters (Lin <-> Two), so it passes unchanged.
71def OptState.map(-P: Type, -Q: Type, st: OptState<P>, f: P -> Q, h: P -> Q) -> OptState<Q>:
72 match st:
73 case NoMom{}: NoMom{}
74 case OneMom{m}: OneMom{f(m)}
75 case TwoMom{m, v}: TwoMom{f(m), h(v)}
76 case Packed{t}: Packed{t}
77
78# A packed tree splits at its node; a part not yet stepped splits into two such parts. A leaf belongs to one matrix and
79# is never split (it goes to the left, the right side starts empty).
80def OptState.split(-A: Type, -B: Type, st: OptState<Two<A, B>>) -> Two<OptState<A>, OptState<B>>:
81 match st:
82 case NoMom{}: Two{NoMom{}, NoMom{}}
83 case OneMom{Two{a, b}}: Two{OneMom{a}, OneMom{b}}
84 case TwoMom{Two{a, b}, Two{c, d}}: Two{TwoMom{a, c}, TwoMom{b, d}}
85 case Packed{PNode{a, b}}: Two{Packed{a}, Packed{b}}
86 case Packed{PNew{+f}}: Two{Packed{PNew{f}}, Packed{PNew{f}}}
87 case Packed{t}: Two{Packed{t}, Packed{PNew{False{}}}}
88
89# Pair the buffers present on both sides. Internal callers split one state, then step both children with the same
90# optimizer, so the variants agree. No buffer is ever allocated to fill an absent side.
91def OptState.join(-A: Type, -B: Type, a: OptState<A>, b: OptState<B>) -> OptState<Two<A, B>>:
92 match a b:
93 case OneMom{x} OneMom{y}: OneMom{Two{x, y}}
94 case TwoMom{x, u} TwoMom{y, v}: TwoMom{Two{x, y}, Two{u, v}}
95 case Packed{x} Packed{y}: Packed{PNode{x, y}}
96 case _ _: NoMom{}
97
98# what an optimizer step returns: the new parameters, the zeroed gradient, the optimizer state
99type Step<-P: Type> is Type:
100 Step{p: P, g: P, st: OptState<P>}
101
102# The optimizer, with the scalars of the current step. Adam's `step` is lr / (1 - b1^t) and `bc2` is
103# sqrt(1 - b2^t), as PyTorch computes them; `Opt.at` fills them from an `OptCfg` and the step count.
104# `dec` picks AdamW's decoupled weight decay (w *= 1 - lr*wd) over Adam's (g += wd*w).
105type Opt is Data:
106 SGD{lr: F32}
107 Momentum{lr: F32, mu: F32}
108 Adam{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, dec: Bool, step: F32, bc2: F32, c1: F32, c2: F32}
109
110# Parameters and their flat dump in a fixed order. Loading has a separate array cursor (PA).
111type PL<-P: Type> is Type:
112 PL{p: P, xs: List<&2, F32>}
113
114# Loading cursor: the shared source array and next element offset, handed back without a list or a clone.
115# The IO boundary checks the source length against Layer.count before calling the pure loaders.
116type PA<-P: Type> is Type:
117 PA{p: P, xs: Array<F32>, off: U32}
118
119type Layer<-P: Type, -X: Type, -Y: Type, -S: Type> is Type:
120 Layer{fwd: P -> X -> Fwd<P, Y, S>, bwd: P -> S -> Y -> P -> Bwd<P, X>, step: Opt -> P -> P -> OptState<P> -> Step<P>, zeros: Unit -> P, dump: P -> PL<P>, load: P -> Array<F32> -> U32 -> PA<P>, store: P -> Array<F32> -> U32 -> PA<P>, count: Unit -> U32}
121
122def Layer.fwd(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>, p: P, x: X) -> Fwd<P, Y, S>:
123 match l:
124 case Layer{fwd, bwd, step, zeros, dump, load, store, count}:
125 fwd(p, x)
126
127def Layer.bwd(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>, p: P, s: S, dy: Y, g: P) -> Bwd<P, X>:
128 match l:
129 case Layer{fwd, bwd, step, zeros, dump, load, store, count}:
130 bwd(p, s, dy, g)
131
132def Layer.step(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>, +o: Opt, p: P, g: P, st: OptState<P>) -> Step<P>:
133 match l:
134 case Layer{fwd, bwd, step, zeros, dump, load, store, count}:
135 step(o, p, g, st)
136
137def Layer.zeros(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>) -> P:
138 match l:
139 case Layer{fwd, bwd, step, zeros, dump, load, store, count}:
140 zeros(Unit{})
141
142# every number of the parameters, in the layer's fixed order (the parameters handed back)
143def Layer.dump(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>, p: P) -> PL<P>:
144 match l:
145 case Layer{fwd, bwd, step, zeros, dump, load, store, count}:
146 dump(p)
147
148# Parameters overwritten from xs at an element offset, in dump order; xs and the next offset come back.
149def Layer.load(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>, p: P, xs: Array<F32>, off: U32) -> PA<P>:
150 match l:
151 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: load(p, xs, off)
152
153# The parameters copied into xs from an element offset, in dump order (load's inverse); xs and the next offset come back.
154def Layer.store(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>, p: P, xs: Array<F32>, off: U32) -> PA<P>:
155 match l:
156 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: store(p, xs, off)
157
158def Layer.count(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer<P, X, Y, S>) -> U32:
159 match l:
160 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: count(Unit{})
161
162# =====================================================================================================
163# Element-wise kernels (templates: one copy per function, no closure at run time)
164# =====================================================================================================
165
166# a[t] = f(a[t]) for the n numbers from t; ra holds the array and a[t], already read
167def ew1_l(~f: @+a: F32 -> F32, nleft: Nat, +t: U32, ra: Array<F32> & F32) -> Array<F32>:
168 match nleft ra:
169 case 0n Tuple{a, v}:
170 a
171 case 1n+p Tuple{a, +v}:
172 ew1_l(~f, p, (t + 1 : U32), Array.get(F32, Array.set(F32, a, t, f(v)), (t + 1 : U32)))
173
174# a[t] = f(a[t], b[t]): the result replaces a; b comes back unchanged
175def ew2_l(~f: @+a: F32 -> @+b: F32 -> F32, nleft: Nat, +t: U32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> TA.P2:
176 match nleft ra rb:
177 case 0n Tuple{a, v} Tuple{b, w}:
178 TA.P2{a, b}
179 case 1n+p Tuple{a, +v} Tuple{b, +w}:
180 ew2_l(~f, p, (t + 1 : U32), Array.get(F32, Array.set(F32, a, t, f(v, w)), (t + 1 : U32)), Array.get(F32, b, (t + 1 : U32)))
181
182def Ew.map(~f: @+a: F32 -> F32, +r: Nat, +c: Nat, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
183 match m:
184 case TA.Mat{d}:
185 TA.Mat{ew1_l(~f, Nat.mul(r, c), 0, Array.get(F32, d, 0))}
186
187def ew2_fin(-r: Nat, -c: Nat, p: TA.P2) -> TA.MPair<r, c>:
188 match p:
189 case TA.P2{a, b}:
190 TA.MPair{TA.Mat{a}, TA.Mat{b}}
191
192# (f(a, b), b), element by element
193def Ew.zip(~f: @+a: F32 -> @+b: F32 -> F32, +r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.MPair<r, c>:
194 match a b:
195 case TA.Mat{x} TA.Mat{y}:
196 ew2_fin(r, c, ew2_l(~f, Nat.mul(r, c), 0, Array.get(F32, x, 0), Array.get(F32, y, 0)))
197
198def f32_add(+a: F32, +b: F32) -> F32:
199 (a + b : F32)
200
201def f32_mul(+a: F32, +b: F32) -> F32:
202 (a * b : F32)
203
204# (a + b, b)
205def Mat.add(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.MPair<r, c>:
206 Ew.zip(~f32_add, r, c, a, b)
207
208def zero_of(+x: F32) -> F32:
209 0.0
210
211# the smallest d with 2^d >= n (TA.cap_depth's answer), by doubling a U32: log2(n) cheap steps, where cap_depth
212# compares Nats one unit at a time at each doubling; f is fuel (32 doublings cover every U32)
213def depth_go(f: Nat, ok: Bool, +n: U32, +p: U32, d: Nat) -> Nat:
214 match f ok:
215 case _ True{}:
216 d
217 case 0n False{}:
218 d
219 case 1n+g False{}:
220 depth_go(g, U32.is_le(n, (p + p : U32)), n, (p + p : U32), 1n+d)
221
222def depth(+n: Nat) -> Nat:
223 depth_go(32n, U32.is_le(U32.from_nat(n), 1), U32.from_nat(n), 1, 0n)
224
225# =====================================================================================================
226# Optimizers: one pass over a parameter, its gradient and the two moment buffers; the gradient is zeroed
227# =====================================================================================================
228
229type Q3 is Type:
230 Q3{w: Array<F32>, g: Array<F32>, m: Array<F32>}
231
232type Q4 is Type:
233 Q4{w: Array<F32>, g: Array<F32>, m: Array<F32>, v: Array<F32>}
234
235type MatStep<-r: Nat, -c: Nat> is Type:
236 MatStep{w: TA.Mat<r, c>, g: TA.Mat<r, c>, st: OptState<TA.Mat<r, c>>}
237
238# w <- w + (-lr)*g, as PyTorch's p.add_(g, alpha=-lr)
239def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array<F32> & F32, rg: Array<F32> & F32) -> TA.P2:
240 match nleft rw rg:
241 case 0n Tuple{w, x} Tuple{g, y}:
242 TA.P2{w, g}
243 case 1n+p Tuple{w, +x} Tuple{g, +y}:
244 sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, Array.set(F32, g, t, 0.0), (t + 1 : U32)))
245
246# buf <- buf*mu + g; w <- w + (-lr)*buf, as PyTorch's SGD with momentum (no dampening, not Nesterov)
247def mom_l(nleft: Nat, +t: U32, +lr: F32, +mu: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<F32> & F32) -> Q3:
248 match nleft rw rg rm:
249 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, b}:
250 Q3{w, g, m}
251 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}:
252 +nb = ((b * mu : F32) + y : F32)
253 +u = (t + 1 : U32)
254 mom_l(p, u, lr, mu, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * nb : F32) : F32)), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nb), u))
255
256# the Adam family, in PyTorch's order of operations (torch.optim.adam._single_tensor_adam):
257# g' = g + wd*w (Adam) or w' = w*(1 - lr*wd) (AdamW); m = m + (1-b1)*(g' - m) (lerp); v = v*b2 + (1-b2)*g'*g';
258# w = w' + (-step*m) / (sqrt(v)/bc2 + eps)
259def adam_grad(dec: Bool, +wd: F32, +w: F32, +g: F32) -> F32:
260 match dec:
261 case True{}:
262 g
263 case False{}:
264 (g + (w * wd : F32) : F32)
265
266def adam_decay(dec: Bool, +lr: F32, +wd: F32, +w: F32) -> F32:
267 match dec:
268 case True{}:
269 (w * (1.0 - (lr * wd : F32) : F32) : F32)
270 case False{}:
271 w
272
273# c1 = 1 - b1 and c2 = 1 - b2 are given as their own constants (0.1, 0.001), as PyTorch computes 1 - beta in float64
274# before rounding: 1 - b2 taken in F32 from b2 = 0.999 (0.99900001287 in F32) is 1.3e-5 too small, and over 100 steps
275# that alone put the tiny Llama 21x further from float64 than PyTorch's float32 (NOTES, 2026-10-18)
276def adam_m(+c1: F32, +m: F32, +g: F32) -> F32:
277 (m + (c1 * (g - m : F32) : F32) : F32)
278
279def adam_v(+b2: F32, +c2: F32, +v: F32, +g: F32) -> F32:
280 ((v * b2 : F32) + ((c2 * g : F32) * g : F32) : F32)
281
282def adam_w(+step: F32, +bc2: F32, +eps: F32, +w: F32, +m: F32, +v: F32) -> F32:
283 (w + ((F32.neg(step) * m : F32) / ((F32.sqrt(v) / bc2 : F32) + eps : F32) : F32) : F32)
284
285def adam_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<F32> & F32, rv: Array<F32> & F32) -> Q4:
286 match nleft rw rg rm rv:
287 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}:
288 Q4{w, g, m, v}
289 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}:
290 +gg = adam_grad(dec, wd, x, y)
291 +nm = adam_m(c1, a, gg)
292 +nv = adam_v(b2, c2, b, gg)
293 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), nm, nv)
294 +u = (t + 1 : U32)
295 adam_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nm), u), Array.get(F32, Array.set(F32, v, t, nv), u))
296
297# ---- AdamW over packed BF16 moments (CAdamWBF16) ----
298# The same F32 arithmetic as adam_l, element by element: each moment is read from its half of a word (BF16 -> F32 is
299# exact), updated in F32, rounded to BF16 (to nearest, ties to even; a NaN becomes 0xFFFF, as torch/dtype.bend and
300# PyTorch's CPU kernel do) and written back into its half; the weight update reads the stored, rounded moments, as
301# PyTorch's AdamW does with BF16 state. The weights and the gradient stay F32.
302
303def ph_f2u(x: F32) -> U32:
304 match x:
305 case F32{w}:
306 U32{w}
307
308def ph_u2f(u: U32) -> F32:
309 match u:
310 case U32{w}:
311 F32{w}
312
313def ph_pick(b: Bool, +x: U32, +y: U32) -> U32:
314 match b:
315 case True{}:
316 x
317 case False{}:
318 y
319
320# F32 -> BF16 bits (dtype.bend's BF16.of_f32): add 0x7FFF plus the lowest kept bit, keep the top half; NaN -> 0xFFFF
321def ph_of_bits(+u: U32) -> U32:
322 ph_pick(U32.is_gt((u .&. 2147483647 : U32), 2139095040), 65535, U32.shrn(((u + 32767 : U32) + (U32.shrn(u, 16n) .&. 1 : U32) : U32), 16n))
323
324def ph_of(x: F32) -> U32:
325 ph_of_bits(ph_f2u(x))
326
327# BF16 bits -> F32, exact
328def ph_to(+h: U32) -> F32:
329 ph_u2f(U32.shln(h, 16n))
330
331# the half of word w that element t uses (odd t: the high half), and w with that half replaced by h
332def ph_half(odd: Bool, +w: U32) -> U32:
333 match odd:
334 case True{}:
335 U32.shrn(w, 16n)
336 case False{}:
337 (w .&. 65535 : U32)
338
339def ph_put(odd: Bool, +w: U32, +h: U32) -> U32:
340 match odd:
341 case True{}:
342 ((w .&. 65535 : U32) .|. U32.shln(h, 16n) : U32)
343 case False{}:
344 ((w .&. 4294901760 : U32) .|. h : U32)
345
346type PQ is Type:
347 PQ{w: Array<F32>, g: Array<F32>, m: Array<U32>, v: Array<U32>}
348
349type PQV is Type:
350 PQV{w: Array<F32>, g: Array<F32>, m: Array<U32>, v: Array<F32>}
351
352# m and v packed. rm and rv hold the word of element t (index t >> 1), read after the previous element's write.
353def adamh_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<U32> & U32, rv: Array<U32> & U32) -> PQ:
354 match nleft rw rg rm rv:
355 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}:
356 PQ{w, g, m, v}
357 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}:
358 +odd = U32.is_eq((t .&. 1 : U32), 1)
359 +gg = adam_grad(dec, wd, x, y)
360 +hm = ph_of(adam_m(c1, ph_to(ph_half(odd, a)), gg))
361 +hv = ph_of(adam_v(b2, c2, ph_to(ph_half(odd, b)), gg))
362 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), ph_to(hm), ph_to(hv))
363 +u = (t + 1 : U32)
364 +k = U32.shrn(t, 1n)
365 +ku = U32.shrn(u, 1n)
366 adamh_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(U32, Array.set(U32, m, k, ph_put(odd, a, hm)), ku), Array.get(U32, Array.set(U32, v, k, ph_put(odd, b, hv)), ku))
367
368# m packed, v in F32 (one cell an element, as adam_l's)
369def adamhv_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array<F32> & F32, rg: Array<F32> & F32, rm: Array<U32> & U32, rv: Array<F32> & F32) -> PQV:
370 match nleft rw rg rm rv:
371 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}:
372 PQV{w, g, m, v}
373 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}:
374 +odd = U32.is_eq((t .&. 1 : U32), 1)
375 +gg = adam_grad(dec, wd, x, y)
376 +hm = ph_of(adam_m(c1, ph_to(ph_half(odd, a)), gg))
377 +nv = adam_v(b2, c2, b, gg)
378 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), ph_to(hm), nv)
379 +u = (t + 1 : U32)
380 adamhv_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(U32, Array.set(U32, m, U32.shrn(t, 1n), ph_put(odd, a, hm)), U32.shrn(u, 1n)), Array.get(F32, Array.set(F32, v, t, nv), u))
381
382# the words of n packed values: ceil(n / 2), in a block of the next power of two
383def ph_words(+n: Nat) -> Array<U32>:
384 Array.new(U32, depth(Nat.div(Nat.add(n, 1n), 2n)), 0)
385
386def ph_fin(-r: Nat, -c: Nat, +n: U32, q: PQ) -> MatStep<r, c>:
387 match q:
388 case PQ{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, Packed{PLeaf{n, m, v}}}
389
390def phv_fin(-r: Nat, -c: Nat, +n: U32, q: PQV) -> MatStep<r, c>:
391 match q:
392 case PQV{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, Packed{PLeafV{n, m, v}}}
393
394# one matrix's step on its leaf; anything else leaves the values untouched, as Opt.mat's fallback
395def adamh_leaf(+r: Nat, +c: Nat, +o: Opt, wa: Array<F32>, ga: Array<F32>, t: PTree) -> MatStep<r, c>:
396 match o t:
397 case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeaf{+n, ma, va}:
398 ph_fin(r, c, n, adamh_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(U32, va, 0)))
399 case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeafV{+n, ma, va}:
400 phv_fin(r, c, n, adamhv_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(F32, va, 0)))
401 case _ t:
402 MatStep{TA.Mat{wa}, TA.Mat{ga}, Packed{t}}
403
404# the leaf of a matrix on its first step: zeros (BF16 0x0000 is +0.0)
405def ph_new(+n: Nat, v16: Bool) -> PTree:
406 match v16:
407 case True{}:
408 PLeaf{U32.from_nat(n), ph_words(n), ph_words(n)}
409 case False{}:
410 PLeafV{U32.from_nat(n), ph_words(n), Array.new(F32, depth(n), 0.0)}
411
412def adamh_mat(+r: Nat, +c: Nat, +o: Opt, wa: Array<F32>, ga: Array<F32>, t: PTree) -> MatStep<r, c>:
413 match t:
414 case PNew{v16}:
415 adamh_leaf(r, c, o, wa, ga, ph_new(Nat.mul(r, c), v16))
416 case t:
417 adamh_leaf(r, c, o, wa, ga, t)
418
419# One matrix step, with only the buffers this optimizer reads.
420def Opt.sgd_fin(-r: Nat, -c: Nat, q: TA.P2) -> MatStep<r, c>:
421 match q:
422 case TA.P2{w, g}: MatStep{TA.Mat{w}, TA.Mat{g}, NoMom{}}
423
424def Opt.mom_fin(-r: Nat, -c: Nat, q: Q3) -> MatStep<r, c>:
425 match q:
426 case Q3{w, g, m}: MatStep{TA.Mat{w}, TA.Mat{g}, OneMom{TA.Mat{m}}}
427
428def Opt.fin(-r: Nat, -c: Nat, q: Q4) -> MatStep<r, c>:
429 match q:
430 case Q4{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, TwoMom{TA.Mat{m}, TA.Mat{v}}}
431
432# A mismatched state is initialized only at the training boundary; Opt.mat requires the matching variant.
433# The fallback leaves the values untouched, so it cannot fabricate a moment buffer.
434def Opt.mat(+r: Nat, +c: Nat, +o: Opt, w: TA.Mat<r, c>, g: TA.Mat<r, c>, st: OptState<TA.Mat<r, c>>) -> MatStep<r, c>:
435 match o w g st:
436 case SGD{lr} TA.Mat{wd} TA.Mat{gd} _:
437 Opt.sgd_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, wd, 0), Array.get(F32, gd, 0)))
438 case Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} OneMom{TA.Mat{md}}:
439 Opt.mom_fin(r, c, mom_l(Nat.mul(r, c), 0, lr, mu, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0)))
440 case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} TwoMom{TA.Mat{md}, TA.Mat{vd}}:
441 Opt.fin(r, c, adam_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd2, dec, step, bc2, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0), Array.get(F32, vd, 0)))
442 case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} Packed{t}:
443 adamh_mat(r, c, o, wd, gd, t)
444 case _ w g st:
445 MatStep{w, g, st}
446
447# What the user picks; `Opt.at` turns it into the step's `Opt` from q1 = 1 - b1^t and q2 = 1 - b2^t (t = 1 on the
448# first step). Adam's c1 and c2 are 1 - b1 and 1 - b2, written as their own constants (adam_m above). CAdamWBF16 is
449# AdamW with its first moment (and its second, when v16) stored as packed BF16 (PTree); weights and gradient stay F32.
450type OptCfg is Data:
451 CSGD{lr: F32}
452 CMomentum{lr: F32, mu: F32}
453 CAdam{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32}
454 CAdamW{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32}
455 CAdamWBF16{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32, v16: Bool}
456
457def Opt.at(+cfg: OptCfg, +q1: F32, +q2: F32) -> Opt:
458 match cfg:
459 case CSGD{lr}:
460 SGD{lr}
461 case CMomentum{lr, mu}:
462 Momentum{lr, mu}
463 case CAdam{lr, b1, b2, eps, wd, c1, c2}:
464 Adam{lr, b1, b2, eps, wd, False{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2}
465 case CAdamW{lr, b1, b2, eps, wd, c1, c2}:
466 Adam{lr, b1, b2, eps, wd, True{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2}
467 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}:
468 Adam{lr, b1, b2, eps, wd, True{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2}
469
470def OptCfg.b1(+cfg: OptCfg) -> F32:
471 match cfg:
472 case CAdam{lr, b1, b2, eps, wd, c1, c2}:
473 b1
474 case CAdamW{lr, b1, b2, eps, wd, c1, c2}:
475 b1
476 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}:
477 b1
478 case _:
479 0.0
480
481def OptCfg.b2(+cfg: OptCfg) -> F32:
482 match cfg:
483 case CAdam{lr, b1, b2, eps, wd, c1, c2}:
484 b2
485 case CAdamW{lr, b1, b2, eps, wd, c1, c2}:
486 b2
487 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}:
488 b2
489 case _:
490 0.0
491
492def OptCfg.c1(+cfg: OptCfg) -> F32:
493 match cfg:
494 case CAdam{lr, b1, b2, eps, wd, c1, c2}:
495 c1
496 case CAdamW{lr, b1, b2, eps, wd, c1, c2}:
497 c1
498 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}:
499 c1
500 case _:
501 0.0
502
503def OptCfg.c2(+cfg: OptCfg) -> F32:
504 match cfg:
505 case CAdam{lr, b1, b2, eps, wd, c1, c2}:
506 c2
507 case CAdamW{lr, b1, b2, eps, wd, c1, c2}:
508 c2
509 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}:
510 c2
511 case _:
512 0.0
513
514# c = 1 - b up to F32's rounding, with 0 < c <= 1 (PyTorch's 0 <= beta < 1). For b in [0, 1), 1 - b taken in F32 is
515# within half an ulp of the exact value and c's own rounding adds at most as much, so a c written as 1 - b is always
516# under 2^-23 from it, while a typo (0.01 for 0.001, c1 and c2 swapped) is far above. A NaN fails every comparison.
517def beta_ok(+b: F32, +c: F32) -> Bool:
518 +d = ((1.0 - b : F32) - c : F32)
519 Bool.and(Bool.and((c > 0.0 : F32), (c <= 1.0 : F32)), ((d * d : F32) <= 0.0000000000000142108547152020037174224853515625 : F32))
520
521# True unless an Adam configuration's c1, c2 disagree with its b1, b2. Train.new poisons the state of a configuration
522# that fails (every loss after the first is NaN); a program checks it first with OptCfg.check to stop with a message.
523def OptCfg.ok(+cfg: OptCfg) -> Bool:
524 match cfg:
525 case CAdam{lr, b1, b2, eps, wd, c1, c2}:
526 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2))
527 case CAdamW{lr, b1, b2, eps, wd, c1, c2}:
528 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2))
529 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}:
530 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2))
531 case _:
532 True{}
533
534def cfg_stop(ok: Bool) -> IO(Unit):
535 match ok:
536 case True{}:
537 IO.pure(Unit, Unit{})
538 case False{}:
539 IO.die(Unit, 1, "CAdam/CAdamW: c1, c2 do not match b1, b2 (c1 = 1 - b1, c2 = 1 - b2, e.g. 0.9, 0.999 with 0.1, 0.001)")
540
541def OptCfg.check(+cfg: OptCfg) -> IO(Unit):
542 cfg_stop(OptCfg.ok(cfg))
543
544# dump and load for a matrix: its r*c numbers, row by row
545def mat_dump_fin(-r: Nat, -c: Nat, ml: TA.ML<r, c>) -> PL<TA.Mat<r, c>>:
546 match ml:
547 case TA.ML{m, xs}:
548 PL{m, xs}
549
550def Mat.dump(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> PL<TA.Mat<r, c>>:
551 mat_dump_fin(r, c, TA.Mat.to_list(r, c, m))
552
553def load_cells(left: Nat, +i: U32, +off: U32, rx: Array<F32> & F32, dst: Array<F32>) -> TA.P2:
554 match left rx:
555 case 0n Tuple{xs, v}: TA.P2{xs, dst}
556 case 1n+q Tuple{xs, v}:
557 load_cells(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, xs, (off + 1 : U32)), Array.set(F32, dst, i, v))
558
559def mat_restore_fin(-r: Nat, -c: Nat, +off: U32, q: TA.P2) -> PA<TA.Mat<r, c>>:
560 match q:
561 case TA.P2{xs, dst}: PA{TA.Mat{dst}, xs, off}
562
563def Mat.restore(+r: Nat, +c: Nat, m: TA.Mat<r, c>, xs: Array<F32>, +off: U32) -> PA<TA.Mat<r, c>>:
564 match m:
565 case TA.Mat{dst}:
566 mat_restore_fin(r, c, (off + U32.from_nat(Nat.mul(r, c)) : U32), load_cells(Nat.mul(r, c), 0, off, Array.get(F32, xs, off), dst))
567
568# load_cells' inverse: copies the matrix's cells i.. into xs from `off`; both arrays come back (the matrix's first)
569def store_cells(left: Nat, +i: U32, +off: U32, rm: Array<F32> & F32, xs: Array<F32>) -> TA.P2:
570 match left rm:
571 case 0n Tuple{ma, v}: TA.P2{ma, xs}
572 case 1n+q Tuple{ma, v}:
573 store_cells(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, ma, (i + 1 : U32)), Array.set(F32, xs, off, v))
574
575def mat_store_fin(-r: Nat, -c: Nat, +off: U32, q: TA.P2) -> PA<TA.Mat<r, c>>:
576 match q:
577 case TA.P2{ma, xs}: PA{TA.Mat{ma}, xs, off}
578
579def Mat.store(+r: Nat, +c: Nat, m: TA.Mat<r, c>, xs: Array<F32>, +off: U32) -> PA<TA.Mat<r, c>>:
580 match m:
581 case TA.Mat{ma}:
582 mat_store_fin(r, c, (off + U32.from_nat(Nat.mul(r, c)) : U32), store_cells(Nat.mul(r, c), 0, off, Array.get(F32, ma, 0), xs))
583
584def pa_second(-A: Type, -B: Type, a: A, rb: PA<B>) -> PA<Two<A, B>>:
585 match rb:
586 case PA{b, xs, off}: PA{Two{a, b}, xs, off}
587
588def pl_pair(-A: Type, -B: Type, ra: PL<A>, rb: PL<B>) -> PL<Two<A, B>>:
589 match ra rb:
590 case PL{a, xa} PL{b, xb}:
591 PL{Two{a, b}, List.append(&2, F32, xa, xb)}
592
593def pl_second(-A: Type, -B: Type, a: A, rb: PL<B>) -> PL<Two<A, B>>:
594 match rb:
595 case PL{b, rest}:
596 PL{Two{a, b}, rest}
597
598# =====================================================================================================
599# Linear: y = x·Wᵀ + b, with W stored out x in (PyTorch's nn.Linear layout); saves its input
600# =====================================================================================================
601
602type Lin<-i: Nat, -o: Nat> is Type:
603 Lin{w: TA.Mat<o, i>, b: TA.Mat<1n, o>}
604
605def lin_f3(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat<o, i>, x: TA.Mat<n, i>, r: TA.MBias<n, o>) -> Fwd<Lin<i, o>, TA.Mat<n, o>, TA.Mat<n, i>>:
606 match r:
607 case TA.MBias{y, b}:
608 Fwd{Lin{w, b}, y, x}
609
610def lin_f2(+n: Nat, -i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: TA.MMulNT<n, i, o>) -> Fwd<Lin<i, o>, TA.Mat<n, o>, TA.Mat<n, i>>:
611 match r:
612 case TA.MMulNT{x, w, z}:
613 lin_f3(n, i, o, w, x, TA.Mat.add_row(n, o, z, b))
614
615def lin_fwd(+n: Nat, +i: Nat, +o: Nat, par: Nat, p: Lin<i, o>, x: TA.Mat<n, i>) -> Fwd<Lin<i, o>, TA.Mat<n, o>, TA.Mat<n, i>>:
616 match p:
617 case Lin{w, b}:
618 lin_f2(n, i, o, b, TA.Mat.matmul_nt(n, i, o, par, x, w))
619
620def lin_b4(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat<o, i>, b: TA.Mat<1n, o>, dx: TA.Mat<n, i>, gw: TA.Mat<o, i>, r: TA.MColSum<n, o>) -> Bwd<Lin<i, o>, TA.Mat<n, i>>:
621 match r:
622 case TA.MColSum{dy, gb}:
623 Bwd{Lin{w, b}, dx, Lin{gw, gb}}
624
625def lin_b3(+n: Nat, -i: Nat, +o: Nat, w: TA.Mat<o, i>, b: TA.Mat<1n, o>, dx: TA.Mat<n, i>, gb: TA.Mat<1n, o>, r: TA.MMulTN<o, n, i>) -> Bwd<Lin<i, o>, TA.Mat<n, i>>:
626 match r:
627 case TA.MMulTN{dy, x, gw}:
628 lin_b4(n, i, o, w, b, dx, gw, TA.Mat.col_sums_acc(n, o, dy, gb))
629
630def lin_b2(+n: Nat, +i: Nat, +o: Nat, par: Nat, b: TA.Mat<1n, o>, x: TA.Mat<n, i>, g: Lin<i, o>, r: TA.MMul<n, o, i>) -> Bwd<Lin<i, o>, TA.Mat<n, i>>:
631 match g r:
632 case Lin{gw, gb} TA.MMul{dy, w, dx}:
633 lin_b3(n, i, o, w, b, dx, gb, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, gw))
634
635# dX = dY·W, dW += dYᵀ·X, db += the column sums of dY; the two sums are added straight into the gradient (the
636# accumulating products of bend-ml-tensor-array 0.1.10.0), with the bits of computing them apart and adding them
637def lin_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: Lin<i, o>, x: TA.Mat<n, i>, dy: TA.Mat<n, o>, g: Lin<i, o>) -> Bwd<Lin<i, o>, TA.Mat<n, i>>:
638 match p:
639 case Lin{w, b}:
640 lin_b2(n, i, o, par, b, x, g, TA.Mat.matmul_precise(n, o, i, par, dy, w))
641
642def pair_s3(-P1: Type, -P2: Type, r1: Step<P1>, r2: Step<P2>) -> Step<Two<P1, P2>>:
643 match r1 r2:
644 case Step{p1, g1, s1} Step{p2, g2, s2}:
645 Step{Two{p1, p2}, Two{g1, g2}, OptState.join(P1, P2, s1, s2)}
646
647def lin_pair(-i: Nat, -o: Nat, p: Lin<i, o>) -> Two<TA.Mat<o, i>, TA.Mat<1n, o>>:
648 match p:
649 case Lin{w, b}: Two{w, b}
650
651def lin_unpair(-i: Nat, -o: Nat, p: Two<TA.Mat<o, i>, TA.Mat<1n, o>>) -> Lin<i, o>:
652 match p:
653 case Two{w, b}: Lin{w, b}
654
655def lin_s2(-i: Nat, -o: Nat, a: MatStep<o, i>, c: MatStep<1n, o>) -> Step<Lin<i, o>>:
656 match a c:
657 case MatStep{w, gw, sw} MatStep{b, gb, sb}:
658 Step{Lin{w, b}, Lin{gw, gb}, OptState.map(Two<TA.Mat<o, i>, TA.Mat<1n, o>>, Lin<i, o>, OptState.join(TA.Mat<o, i>, TA.Mat<1n, o>, sw, sb), x => lin_unpair(i, o, x), x => lin_unpair(i, o, x))}
659
660def lin_step2(+i: Nat, +o: Nat, +op: Opt, p: Lin<i, o>, g: Lin<i, o>, st: Two<OptState<TA.Mat<o, i>>, OptState<TA.Mat<1n, o>>>) -> Step<Lin<i, o>>:
661 match p g st:
662 case Lin{w, b} Lin{gw, gb} Two{sw, sb}:
663 lin_s2(i, o, Opt.mat(o, i, op, w, gw, sw), Opt.mat(1n, o, op, b, gb, sb))
664
665def lin_step(+i: Nat, +o: Nat, +op: Opt, p: Lin<i, o>, g: Lin<i, o>, st: OptState<Lin<i, o>>) -> Step<Lin<i, o>>:
666 lin_step2(i, o, op, p, g, OptState.split(TA.Mat<o, i>, TA.Mat<1n, o>, OptState.map(Lin<i, o>, Two<TA.Mat<o, i>, TA.Mat<1n, o>>, st, x => lin_pair(i, o, x), x => lin_pair(i, o, x))))
667
668def lin_dump2(-i: Nat, -o: Nat, r: PL<Two<TA.Mat<o, i>, TA.Mat<1n, o>>>) -> PL<Lin<i, o>>:
669 match r:
670 case PL{Two{w, b}, xs}:
671 PL{Lin{w, b}, xs}
672
673def lin_dump(+i: Nat, +o: Nat, p: Lin<i, o>) -> PL<Lin<i, o>>:
674 match p:
675 case Lin{w, b}:
676 lin_dump2(i, o, pl_pair(TA.Mat<o, i>, TA.Mat<1n, o>, Mat.dump(o, i, w), Mat.dump(1n, o, b)))
677
678def lin_load3(-i: Nat, -o: Nat, r: PA<Two<TA.Mat<o, i>, TA.Mat<1n, o>>>) -> PA<Lin<i, o>>:
679 match r:
680 case PA{Two{w, b}, xs, off}: PA{Lin{w, b}, xs, off}
681
682def lin_load2(+i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: PA<TA.Mat<o, i>>) -> PA<Lin<i, o>>:
683 match r:
684 case PA{w, xs, off}:
685 lin_load3(i, o, pa_second(TA.Mat<o, i>, TA.Mat<1n, o>, w, Mat.restore(1n, o, b, xs, off)))
686
687def lin_load(+i: Nat, +o: Nat, p: Lin<i, o>, xs: Array<F32>, off: U32) -> PA<Lin<i, o>>:
688 match p:
689 case Lin{w, b}: lin_load2(i, o, b, Mat.restore(o, i, w, xs, off))
690
691def lin_store2(+i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: PA<TA.Mat<o, i>>) -> PA<Lin<i, o>>:
692 match r:
693 case PA{w, xs, off}:
694 lin_load3(i, o, pa_second(TA.Mat<o, i>, TA.Mat<1n, o>, w, Mat.store(1n, o, b, xs, off)))
695
696def lin_store(+i: Nat, +o: Nat, p: Lin<i, o>, xs: Array<F32>, off: U32) -> PA<Lin<i, o>>:
697 match p:
698 case Lin{w, b}: lin_store2(i, o, b, Mat.store(o, i, w, xs, off))
699
700def lin_zeros(+i: Nat, +o: Nat) -> Lin<i, o>:
701 Lin{TA.Mat.zeros(o, i), TA.Mat.zeros(1n, o)}
702
703# a linear layer from i to o features on batches of n rows; its products run in 2^par blocks of rows
704def Linear(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer<Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
705 Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => lin_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => lin_zeros(i, o), p => lin_dump(i, o, p), p => xs => off => lin_load(i, o, p, xs, off), p => xs => off => lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))}
706
707# the first layer of a model, whose input (the data) needs no gradient, as an input without requires_grad in PyTorch:
708# the backward pass computes dW and db only, skipping dX = dY·W (a third of a linear layer's work), and returns zeros
709# as dX. Use it only where nothing before the layer has parameters.
710def linf_b3(+n: Nat, +i: Nat, +o: Nat, w: TA.Mat<o, i>, b: TA.Mat<1n, o>, gb: TA.Mat<1n, o>, r: TA.MMulTN<o, n, i>) -> Bwd<Lin<i, o>, TA.Mat<n, i>>:
711 match r:
712 case TA.MMulTN{dy, x, gw}:
713 lin_b4(n, i, o, w, b, TA.Mat.zeros(n, i), gw, TA.Mat.col_sums_acc(n, o, dy, gb))
714
715def linf_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: Lin<i, o>, x: TA.Mat<n, i>, dy: TA.Mat<n, o>, g: Lin<i, o>) -> Bwd<Lin<i, o>, TA.Mat<n, i>>:
716 match p g:
717 case Lin{w, b} Lin{gw, gb}:
718 linf_b3(n, i, o, w, b, gb, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, gw))
719
720def LinearFirst(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer<Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
721 Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => linf_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => lin_zeros(i, o), p => lin_dump(i, o, p), p => xs => off => lin_load(i, o, p, xs, off), p => xs => off => lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))}
722
723# =====================================================================================================
724# Layers without parameters
725# =====================================================================================================
726
727def pick(b: Bool, +x: F32, +y: F32) -> F32:
728 match b:
729 case True{}:
730 x
731 case False{}:
732 y
733
734def no_step(-P: Type, +op: Opt, p: P, g: P, st: OptState<P>) -> Step<P>:
735 Step{p, g, st}
736
737# an element-wise activation that saves a copy of its output (relu, sigmoid, tanh: their derivative is a function of
738# the output). `f` is the activation, `d(dy, y)` the gradient of its input.
739def act_f2(-n: Nat, -c: Nat, r: TA.MM2<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
740 match r:
741 case TA.MM2{y, s}:
742 Fwd{Unit{}, y, s}
743
744def act_b2(-n: Nat, -c: Nat, r: TA.MPair<n, c>) -> Bwd<Unit, TA.Mat<n, c>>:
745 match r:
746 case TA.MPair{dx, y}:
747 Bwd{Unit{}, dx, Unit{}}
748
749def ActOut(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
750 Layer{p => x => act_f2(n, c, TA.Mat.clone(n, c, Ew.map(~f, n, c, x))), p => s => dy => g => act_b2(n, c, Ew.zip(~d, n, c, dy, s)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
751
752# relu(x) = x unless x < 0 (so relu(NaN) = NaN and relu(-0) = -0); its gradient is +0 where the output is <= 0 and
753# passes everywhere else, NaN included (PyTorch's threshold_backward)
754def relu_v(+x: F32) -> F32:
755 pick(F32.is_lt(x, 0.0), 0.0, x)
756
757def relu_d(+dy: F32, +y: F32) -> F32:
758 pick(F32.is_le(y, 0.0), 0.0, dy)
759
760def ReLU(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
761 ActOut(~relu_v, ~relu_d, n, c)
762
763# =====================================================================================================
764# Seq: la then lb. Parameters, gradients, saved values and optimizer state are pairs of the two layers'.
765# =====================================================================================================
766
767def seq_f3(-P1: Type, -P2: Type, -Z: Type, -S1: Type, -S2: Type, p1: P1, s1: S1, r: Fwd<P2, Z, S2>) -> Fwd<Two<P1, P2>, Z, Two<S1, S2>>:
768 match r:
769 case Fwd{p2, z, s2}:
770 Fwd{Two{p1, p2}, z, Two{s1, s2}}
771
772def seq_f2(~P2: Type, ~Y: Type, ~Z: Type, ~S2: Type, ~lb: Layer<P2, Y, Z, S2>, -P1: Type, -S1: Type, p2: P2, r: Fwd<P1, Y, S1>) -> Fwd<Two<P1, P2>, Z, Two<S1, S2>>:
773 match r:
774 case Fwd{p1, y, s1}:
775 seq_f3(P1, P2, Z, S1, S2, p1, s1, Layer.fwd(P2, Y, Z, S2, lb, p2, y))
776
777def seq_fwd(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X, Y, S1>, ~lb: Layer<P2, Y, Z, S2>, p: Two<P1, P2>, x: X) -> Fwd<Two<P1, P2>, Z, Two<S1, S2>>:
778 match p:
779 case Two{p1, p2}:
780 seq_f2(~P2, ~Y, ~Z, ~S2, ~lb, P1, S1, p2, Layer.fwd(P1, X, Y, S1, la, p1, x))
781
782def seq_b3(-P1: Type, -P2: Type, -X: Type, p2: P2, g2: P2, r: Bwd<P1, X>) -> Bwd<Two<P1, P2>, X>:
783 match r:
784 case Bwd{p1, dx, g1}:
785 Bwd{Two{p1, p2}, dx, Two{g1, g2}}
786
787def seq_b2(~P1: Type, ~X: Type, ~Y: Type, ~S1: Type, ~la: Layer<P1, X, Y, S1>, -P2: Type, p1: P1, s1: S1, g1: P1, r: Bwd<P2, Y>) -> Bwd<Two<P1, P2>, X>:
788 match r:
789 case Bwd{p2, dy, g2}:
790 seq_b3(P1, P2, X, p2, g2, Layer.bwd(P1, X, Y, S1, la, p1, s1, dy, g1))
791
792def seq_bwd(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X, Y, S1>, ~lb: Layer<P2, Y, Z, S2>, p: Two<P1, P2>, s: Two<S1, S2>, dz: Z, g: Two<P1, P2>) -> Bwd<Two<P1, P2>, X>:
793 match p s g:
794 case Two{p1, p2} Two{s1, s2} Two{g1, g2}:
795 seq_b2(~P1, ~X, ~Y, ~S1, ~la, P2, p1, s1, g1, Layer.bwd(P2, Y, Z, S2, lb, p2, s2, dz, g2))
796
797def seq_s3(-P1: Type, -P2: Type, r1: Step<P1>, r2: Step<P2>) -> Step<Two<P1, P2>>:
798 pair_s3(P1, P2, r1, r2)
799
800def seq_s2(~P2: Type, ~Y: Type, ~Z: Type, ~S2: Type, ~lb: Layer<P2, Y, Z, S2>, -P1: Type, +op: Opt, p2: P2, g2: P2, s2: OptState<P2>, r1: Step<P1>) -> Step<Two<P1, P2>>:
801 seq_s3(P1, P2, r1, Layer.step(P2, Y, Z, S2, lb, op, p2, g2, s2))
802
803def seq_step2(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X, Y, S1>, ~lb: Layer<P2, Y, Z, S2>, +op: Opt, p: Two<P1, P2>, g: Two<P1, P2>, st: Two<OptState<P1>, OptState<P2>>) -> Step<Two<P1, P2>>:
804 match p g st:
805 case Two{p1, p2} Two{g1, g2} Two{s1, s2}:
806 seq_s2(~P2, ~Y, ~Z, ~S2, ~lb, P1, op, p2, g2, s2, Layer.step(P1, X, Y, S1, la, op, p1, g1, s1))
807
808def seq_step(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X, Y, S1>, ~lb: Layer<P2, Y, Z, S2>, +op: Opt, p: Two<P1, P2>, g: Two<P1, P2>, st: OptState<Two<P1, P2>>) -> Step<Two<P1, P2>>:
809 seq_step2(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, op, p, g, OptState.split(P1, P2, st))
810
811# dump and load for a pair of layers: the first layer's numbers, then the second's
812def pair_dump(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, p: Two<P1, P2>) -> PL<Two<P1, P2>>:
813 match p:
814 case Two{p1, p2}:
815 pl_pair(P1, P2, Layer.dump(P1, X1, Y1, S1, la, p1), Layer.dump(P2, X2, Y2, S2, lb, p2))
816
817def pair_load2(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer<P2, X2, Y2, S2>, -P1: Type, p2: P2, r: PA<P1>) -> PA<Two<P1, P2>>:
818 match r:
819 case PA{p1, xs, off}:
820 pa_second(P1, P2, p1, Layer.load(P2, X2, Y2, S2, lb, p2, xs, off))
821
822def pair_load(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, p: Two<P1, P2>, xs: Array<F32>, off: U32) -> PA<Two<P1, P2>>:
823 match p:
824 case Two{p1, p2}:
825 pair_load2(~P2, ~X2, ~Y2, ~S2, ~lb, P1, p2, Layer.load(P1, X1, Y1, S1, la, p1, xs, off))
826
827def pair_store2(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer<P2, X2, Y2, S2>, -P1: Type, p2: P2, r: PA<P1>) -> PA<Two<P1, P2>>:
828 match r:
829 case PA{p1, xs, off}:
830 pa_second(P1, P2, p1, Layer.store(P2, X2, Y2, S2, lb, p2, xs, off))
831
832def pair_store(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, p: Two<P1, P2>, xs: Array<F32>, off: U32) -> PA<Two<P1, P2>>:
833 match p:
834 case Two{p1, p2}:
835 pair_store2(~P2, ~X2, ~Y2, ~S2, ~lb, P1, p2, Layer.store(P1, X1, Y1, S1, la, p1, xs, off))
836
837def Seq(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X, Y, S1>, ~lb: Layer<P2, Y, Z, S2>) -> Layer<Two<P1, P2>, X, Z, Two<S1, S2>>:
838 Layer{p => x => seq_fwd(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, p, x), p => s => dz => g => seq_bwd(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, p, s, dz, g), op => p => g => st => seq_step(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P1, X, Y, S1, la), Layer.zeros(P2, Y, Z, S2, lb)}, p => pair_dump(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X, Y, S1, la) + Layer.count(P2, Y, Z, S2, lb) : U32)}
839
840# Owned layer composition for dimensions read at run time. Each child record is opened once; each closure goes to
841# exactly one corresponding closure of the composite layer. The template-based Seq above keeps its existing laws.
842def seq_at_f2(-P1: Type, -P2: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, p2: P2, fb: P2 -> Y -> Fwd<P2, Z, S2>, r: Fwd<P1, Y, S1>) -> Fwd<Two<P1, P2>, Z, Two<S1, S2>>:
843 match r:
844 case Fwd{p1, y, s1}: seq_f3(P1, P2, Z, S1, S2, p1, s1, fb(p2, y))
845
846def seq_at_f(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, fa: P1 -> X -> Fwd<P1, Y, S1>, fb: P2 -> Y -> Fwd<P2, Z, S2>, p: Two<P1, P2>, x: X) -> Fwd<Two<P1, P2>, Z, Two<S1, S2>>:
847 match p:
848 case Two{p1, p2}: seq_at_f2(P1, P2, Y, Z, S1, S2, p2, fb, fa(p1, x))
849
850def seq_at_b2(-P1: Type, -P2: Type, -X: Type, -Y: Type, -S1: Type, ba: P1 -> S1 -> Y -> P1 -> Bwd<P1, X>, p1: P1, s1: S1, g1: P1, r: Bwd<P2, Y>) -> Bwd<Two<P1, P2>, X>:
851 match r:
852 case Bwd{p2, dy, g2}: seq_b3(P1, P2, X, p2, g2, ba(p1, s1, dy, g1))
853
854def seq_at_b(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, ba: P1 -> S1 -> Y -> P1 -> Bwd<P1, X>, bb: P2 -> S2 -> Z -> P2 -> Bwd<P2, Y>, p: Two<P1, P2>, s: Two<S1, S2>, dz: Z, g: Two<P1, P2>) -> Bwd<Two<P1, P2>, X>:
855 match p s g:
856 case Two{p1, p2} Two{s1, s2} Two{g1, g2}: seq_at_b2(P1, P2, X, Y, S1, ba, p1, s1, g1, bb(p2, s2, dz, g2))
857
858def seq_at_step(-P1: Type, -P2: Type, sa: Opt -> P1 -> P1 -> OptState<P1> -> Step<P1>, sb: Opt -> P2 -> P2 -> OptState<P2> -> Step<P2>, +op: Opt, p: Two<P1, P2>, g: Two<P1, P2>, st: Two<OptState<P1>, OptState<P2>>) -> Step<Two<P1, P2>>:
859 match p g st:
860 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: pair_s3(P1, P2, sa(op, p1, g1, s1), sb(op, p2, g2, s2))
861
862def seq_at_dump(-P1: Type, -P2: Type, da: P1 -> PL<P1>, db: P2 -> PL<P2>, p: Two<P1, P2>) -> PL<Two<P1, P2>>:
863 match p:
864 case Two{p1, p2}: pl_pair(P1, P2, da(p1), db(p2))
865
866def seq_at_load2(-P1: Type, -P2: Type, lb: P2 -> Array<F32> -> U32 -> PA<P2>, p2: P2, r: PA<P1>) -> PA<Two<P1, P2>>:
867 match r:
868 case PA{p1, xs, off}: pa_second(P1, P2, p1, lb(p2, xs, off))
869
870def seq_at_load(-P1: Type, -P2: Type, la: P1 -> Array<F32> -> U32 -> PA<P1>, lb: P2 -> Array<F32> -> U32 -> PA<P2>, p: Two<P1, P2>, xs: Array<F32>, off: U32) -> PA<Two<P1, P2>>:
871 match p:
872 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off))
873
874def seq_at_store(-P1: Type, -P2: Type, la: P1 -> Array<F32> -> U32 -> PA<P1>, lb: P2 -> Array<F32> -> U32 -> PA<P2>, p: Two<P1, P2>, xs: Array<F32>, off: U32) -> PA<Two<P1, P2>>:
875 match p:
876 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off))
877
878def Seq.at(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, a: Layer<P1, X, Y, S1>, b: Layer<P2, Y, Z, S2>) -> Layer<Two<P1, P2>, X, Z, Two<S1, S2>>:
879 match a b:
880 case Layer{fa, ba, sa, za, da, la, ta, na} Layer{fb, bb, sb, zb, db, lb, tb, nb}:
881 Layer{p => x => seq_at_f(P1, P2, X, Y, Z, S1, S2, fa, fb, p, x), p => s => dz => g => seq_at_b(P1, P2, X, Y, Z, S1, S2, ba, bb, p, s, dz, g), op => p => g => st => seq_at_step(P1, P2, sa, sb, op, p, g, OptState.split(P1, P2, st)), u => Two{za(Unit{}), zb(Unit{})}, p => seq_at_dump(P1, P2, da, db, p), p => xs => off => seq_at_load(P1, P2, la, lb, p, xs, off), p => xs => off => seq_at_store(P1, P2, ta, tb, p, xs, off), u => (na(Unit{}) + nb(Unit{}) : U32)}
882
883# =====================================================================================================
884# Losses: the mean loss over the batch, and its gradient with respect to the output
885# =====================================================================================================
886
887type LossOut<-Y: Type> is Type:
888 LossOut{loss: F32, dy: Y}
889
890# =====================================================================================================
891# Training: one step of any layer under any loss
892# =====================================================================================================
893
894# Parameters, gradient (zero between steps), optimizer-specific buffers, Adam's 1 - b1^t and 1 - b2^t, last loss.
895type Train<-P: Type> is Type:
896 Train{p: P, g: P, st: OptState<P>, p1: F32, p2: F32, loss: F32}
897
898def OptState.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, cfg: OptCfg) -> OptState<P>:
899 match cfg:
900 case CSGD{lr}: NoMom{}
901 case CMomentum{lr, mu}: OneMom{Layer.zeros(P, X, Y, S, m)}
902 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: Packed{PNew{v16}}
903 case _: TwoMom{Layer.zeros(P, X, Y, S, m), Layer.zeros(P, X, Y, S, m)}
904
905# 1 - b^t starts at 0; a configuration OptCfg.ok rejects starts at NaN, so every loss after the first is NaN.
906def q_start(ok: Bool) -> F32:
907 match ok:
908 case True{}:
909 0.0
910 case False{}:
911 (0.0 / 0.0 : F32)
912
913def Train.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +cfg: OptCfg, p: P) -> Train<P>:
914 +q = q_start(OptCfg.ok(cfg))
915 Train{p, Layer.zeros(P, X, Y, S, m), OptState.new(~P, ~X, ~Y, ~S, ~m, cfg), q, q, 0.0}
916
917def ts5(-P: Type, +l: F32, +p1: F32, +p2: F32, r: Step<P>) -> Train<P>:
918 match r:
919 case Step{p, g, st}:
920 Train{p, g, st, p1, p2, l}
921
922def ts4(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +cfg: OptCfg, st: OptState<P>, +p1: F32, +p2: F32, +l: F32, r: Bwd<P, X>) -> Train<P>:
923 match r:
924 case Bwd{p, dx, g}:
925 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32)
926 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32)
927 ts5(P, l, q1, q2, Layer.step(P, X, Y, S, m, Opt.at(cfg, q1, q2), p, g, st))
928
929def ts3(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +cfg: OptCfg, p: P, s: S, g: P, st: OptState<P>, +p1: F32, +p2: F32, r: LossOut<Y>) -> Train<P>:
930 match r:
931 case LossOut{l, dy}:
932 ts4(~P, ~X, ~Y, ~S, ~m, cfg, st, p1, p2, l, Layer.bwd(P, X, Y, S, m, p, s, dy, g))
933
934def ts2(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~T: Type, ~m: Layer<P, X, Y, S>, ~loss: Y -> T -> LossOut<Y>, +cfg: OptCfg, g: P, st: OptState<P>, +p1: F32, +p2: F32, t: T, r: Fwd<P, Y, S>) -> Train<P>:
935 match r:
936 case Fwd{p, y, s}:
937 ts3(~P, ~X, ~Y, ~S, ~m, cfg, p, s, g, st, p1, p2, loss(y, t))
938
939# forward, loss, backward and one optimizer step: the same code for every model
940def Train.step(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~T: Type, ~m: Layer<P, X, Y, S>, ~loss: Y -> T -> LossOut<Y>, +cfg: OptCfg, tr: Train<P>, x: X, t: T) -> Train<P>:
941 match tr:
942 case Train{p, g, st, p1, p2, l}:
943 ts2(~P, ~X, ~Y, ~S, ~T, ~m, ~loss, cfg, g, st, p1, p2, t, Layer.fwd(P, X, Y, S, m, p, x))
944
945# One owned model, including dimensions captured at run time. Its fwd, bwd and step closures each run once.
946def train_at_b(-P: Type, -X: Type, sf: Opt -> P -> P -> OptState<P> -> Step<P>, +cfg: OptCfg, st: OptState<P>, +p1: F32, +p2: F32, +l: F32, r: Bwd<P, X>) -> Train<P>:
947 match r:
948 case Bwd{p, dx, g}:
949 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32)
950 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32)
951 ts5(P, l, q1, q2, sf(Opt.at(cfg, q1, q2), p, g, st))
952
953def train_at_loss(-P: Type, -X: Type, -Y: Type, -S: Type, bf: P -> S -> Y -> P -> Bwd<P, X>, sf: Opt -> P -> P -> OptState<P> -> Step<P>, +cfg: OptCfg, p: P, s: S, g: P, st: OptState<P>, +p1: F32, +p2: F32, r: LossOut<Y>) -> Train<P>:
954 match r:
955 case LossOut{l, dy}: train_at_b(P, X, sf, cfg, st, p1, p2, l, bf(p, s, dy, g))
956
957def train_at_f(-P: Type, -X: Type, -Y: Type, -S: Type, -Target: Type, bf: P -> S -> Y -> P -> Bwd<P, X>, sf: Opt -> P -> P -> OptState<P> -> Step<P>, loss: Y -> Target -> LossOut<Y>, +cfg: OptCfg, g: P, st: OptState<P>, +p1: F32, +p2: F32, t: Target, r: Fwd<P, Y, S>) -> Train<P>:
958 match r:
959 case Fwd{p, y, s}: train_at_loss(P, X, Y, S, bf, sf, cfg, p, s, g, st, p1, p2, loss(y, t))
960
961def Train.step_at(-P: Type, -X: Type, -Y: Type, -S: Type, -Target: Type, m: Layer<P, X, Y, S>, loss: Y -> Target -> LossOut<Y>, +cfg: OptCfg, tr: Train<P>, x: X, t: Target) -> Train<P>:
962 match m tr:
963 case Layer{ff, bf, sf, zeros, dump, load, store, count} Train{p, g, st, p1, p2, l}:
964 train_at_f(P, X, Y, S, Target, bf, sf, loss, cfg, g, st, p1, p2, t, ff(p, x))
965
966# =====================================================================================================
967# Reading and printing numbers (bend-ml's and llama.bend's loaders, MIT)
968# =====================================================================================================
969
970# U32 and F32 hold the same Word(32n): these are bitcasts
971def bitcast(u: U32) -> F32:
972 match u:
973 case U32{w}:
974 F32{w}
975
976def bits(x: F32) -> U32:
977 match x:
978 case F32{w}:
979 U32{w}
980
981def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32:
982 bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32))
983
984def f32s(bs: List<&2, U32>) -> List<&2, F32>:
985 match bs:
986 case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}:
987 f32_le(b0, b1, b2, b3) <> f32s(t)
988 case _:
989 Nil{}
990
991def rd_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>):
992 match m:
993 case Tuple{f, r}:
994 do IO<List<&2, U32>>:
995 bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
996 File.close(f)
997 return bs
998
999# n bytes from offset off
1000def read_at(path: String, +off: U32, +n: U32) -> IO(List<&2, U32>):
1001 do IO<List<&2, U32>>:
1002 f : File <- IO.try(File, File.open(path, "r"))
1003 m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, off, n)
1004 rd_done(m)
1005
1006def mat_or_die(-r: Nat, -c: Nat, +path: String, m: Maybe<&1, TA.Mat<r, c>>) -> IO(TA.Mat<r, c>):
1007 match m:
1008 case None{}:
1009 IO.die(TA.Mat<r, c>, 1, path ++ ": wrong number of values for this shape")
1010 case Some{x}:
1011 IO.pure(TA.Mat<r, c>, x)
1012
1013# an r x c matrix stored as r*c little-endian F32 from byte off of a file
1014def Mat.load(+r: Nat, +c: Nat, +path: String, +off: U32) -> IO(TA.Mat<r, c>):
1015 do IO<TA.Mat<r, c>>:
1016 bs : List<&2, U32> <- read_at(path, off, (U32.from_nat(Nat.mul(r, c)) * 4 : U32))
1017 mat_or_die(r, c, path, TA.Mat.of_list(r, c, f32s(bs)))
1018
1019def words(xs: List<&2, F32>) -> String:
1020 match xs:
1021 case Nil{}:
1022 ""
1023 case Con{h, t}:
1024 " " ++ U32.show(bits(h)) ++ words(t)
1025
1026def ml_words(-r: Nat, -c: Nat, m: TA.ML<r, c>) -> String:
1027 match m:
1028 case TA.ML{d, xs}:
1029 words(xs)
1030
1031# one line: the name, then each number as the decimal of its 32 bits (exact, for the tests in reference/)
1032def Mat.print(+r: Nat, +c: Nat, +name: String, m: TA.Mat<r, c>) -> IO(Unit):
1033 IO.print(name ++ ml_words(r, c, TA.Mat.to_list(r, c, m)))
1034
1035def F32.print(+name: String, x: F32) -> IO(Unit):
1036 IO.print(name ++ " " ++ U32.show(bits(x)))
1037
1038# n zeros, in an array with room for them
1039def zeros_a(+n: Nat) -> Array<F32>:
1040 Array.new(F32, depth(n), 0.0)
1041
1042# =====================================================================================================
1043# More element-wise activations
1044# =====================================================================================================
1045
1046# an activation whose derivative needs its input (silu, gelu): saves a copy of the input
1047def acti_f3(~f: @+a: F32 -> F32, +n: Nat, +c: Nat, m: TA.MM2<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1048 match m:
1049 case TA.MM2{x, s}:
1050 Fwd{Unit{}, Ew.map(~f, n, c, x), s}
1051
1052def ActIn(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1053 Layer{p => x => acti_f3(~f, n, c, TA.Mat.clone(n, c, x)), p => s => dy => g => act_b2(n, c, Ew.zip(~d, n, c, dy, s)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
1054
1055def sigmoid_v(+x: F32) -> F32:
1056 (1.0 / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
1057
1058# PyTorch's sigmoid_backward: dy * (1 - y) * y
1059def sigmoid_d(+dy: F32, +y: F32) -> F32:
1060 ((dy * (1.0 - y : F32) : F32) * y : F32)
1061
1062def Sigmoid(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1063 ActOut(~sigmoid_v, ~sigmoid_d, n, c)
1064
1065def tanh_v(+x: F32) -> F32:
1066 F32.tanh(x)
1067
1068def tanh_d(+dy: F32, +y: F32) -> F32:
1069 (dy * (1.0 - (y * y : F32) : F32) : F32)
1070
1071def Tanh(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1072 ActOut(~tanh_v, ~tanh_d, n, c)
1073
1074# silu(x) = x * sigmoid(x), computed as PyTorch's CPU kernel: x / (1 + exp(-x));
1075# its gradient dy * s * (1 + x * (1 - s)) with s = sigmoid(x)
1076def silu_v(+x: F32) -> F32:
1077 (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
1078
1079def silu_d(+dy: F32, +x: F32) -> F32:
1080 +s = sigmoid_v(x)
1081 ((dy * s : F32) * (1.0 + (x * (1.0 - s : F32) : F32) : F32) : F32)
1082
1083def SiLU(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1084 ActIn(~silu_v, ~silu_d, n, c)
1085
1086# GELU in its tanh form (PyTorch's approximate="tanh"): 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3))).
1087# The exact form, with erf, is GELUErf below.
1088def gelu_inner(+x: F32) -> F32:
1089 (0.7978845608028654 * (x + (0.044715 * ((x * x : F32) * x : F32) : F32) : F32) : F32)
1090
1091def gelu_v(+x: F32) -> F32:
1092 (((0.5 * x : F32) * (1.0 + F32.tanh(gelu_inner(x)) : F32)) : F32)
1093
1094def gelu_d(+dy: F32, +x: F32) -> F32:
1095 +th = F32.tanh(gelu_inner(x))
1096 +left = (0.5 * x : F32)
1097 +dleft = (0.5 * (1.0 + th : F32) : F32)
1098 +dright = ((left * (1.0 - (th * th : F32) : F32) : F32) * (0.7978845608028654 * (1.0 + ((3.0 * 0.044715 : F32) * (x * x : F32) : F32) : F32) : F32) : F32)
1099 (dy * (dleft + dright : F32) : F32)
1100
1101def GELU(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1102 ActIn(~gelu_v, ~gelu_d, n, c)
1103
1104# =====================================================================================================
1105# erf in plain Bend, and GELU in its exact form
1106# =====================================================================================================
1107
1108# erf(x) = 2/sqrt(pi) times the integral of exp(-t^2) from 0 to x. Base has no erf, so it is written here with F32
1109# additions, multiplications and comparisons only: no libm call, so every target runs the same operations (the emitted C
1110# turns fused multiply-adds off); the C and JavaScript programs give the same bits (CUDA is not checked yet). erf is
1111# odd, erf(-x) = -erf(x), so the work is done on a = |x| and the sign put back, in three kinds of pieces:
1112# - a < 0.84375: erf(x) = x + x P(x^2), with P of degree 6 (fdlibm's erff has this x + x y shape there); it takes
1113# x itself, so erf(-0) = -0, and NaN reaches it (every comparison with NaN is false) and comes out NaN;
1114# - 0.84375 <= a < 3.92, in six intervals around c = 1, 1.5, ..., 3.5: erf(a) = K + S(a - c), where K is erf(c)
1115# rounded to F32 and S a polynomial of degree 7 or 8 whose value is small next to K, so the last addition rounds
1116# once and carries almost all of the error (a - c is exact: a and c are within a factor of 2);
1117# - a >= 3.92, infinity included: 1, since erf(a) rounds to 1 in F32 from a = 3.9192 on.
1118# The coefficients are least-squares fits against a 50-digit erf on Chebyshev points, rounded to F32. Measured on all
1119# 2^32 inputs against glibc's double erf: at most 0.93 ulp off, the correctly rounded F32 on 97.1% of the inputs, and 1
1120# ulp from it on the rest; NaN for every NaN.
1121
1122# one Horner step, c + t s; erf_hm is the step for a negative coefficient -c: t s - c gives the bits of -c + t s
1123def erf_hp(+c: F32, +t: F32, +s: F32) -> F32:
1124 (c + (t * s : F32) : F32)
1125
1126def erf_hm(+c: F32, +t: F32, +s: F32) -> F32:
1127 ((t * s : F32) - c : F32)
1128
1129def erf_small(+x: F32) -> F32:
1130 +z = (x * x : F32)
1131 (x + (x * erf_hp(0.12837916612625122, z, erf_hm(0.37612634897232056, z, erf_hp(0.11283732205629349, z, erf_hm(0.026861993595957756, z, erf_hp(0.005208558402955532, z, erf_hm(0.000823960464913398, z, 8.829222497297451e-05)))))) : F32) : F32)
1132
1133# [0.84375, 1.25): t = a - 1.0
1134def erf_p0(+t: F32) -> F32:
1135 (0.8427007794380188 + erf_hp(1.3508899776581984e-08, t, erf_hp(0.41510748863220215, t, erf_hm(0.41510751843452454, t, erf_hp(0.13837093114852905, t, erf_hp(0.06918490678071976, t, erf_hm(0.06929104775190353, t, erf_hp(0.004778923466801643, t, erf_hp(0.017144912853837013, t, F32.neg(0.010648881085216999))))))))) : F32)
1136
1137# [1.25, 1.75): t = a - 1.5
1138def erf_p1(+t: F32) -> F32:
1139 (0.9661051630973816 + erf_hm(1.6622220044837377e-08, t, erf_hp(0.11893028765916824, t, erf_hm(0.17839543521404266, t, erf_hp(0.1387522667646408, t, erf_hm(0.04459875077009201, t, erf_hm(0.014879402704536915, t, erf_hp(0.019323812797665596, t, erf_hm(0.004488371312618256, t, F32.neg(0.0023522074334323406))))))))) : F32)
1140
1141# [1.75, 2.25): t = a - 2.0
1142def erf_p2(+t: F32) -> F32:
1143 (0.9953222870826721 + erf_hm(2.2064206461891445e-08, t, erf_hp(0.02066698484122753, t, erf_hm(0.04133396968245506, t, erf_hp(0.04822301119565964, t, erf_hm(0.03444507718086243, t, erf_hp(0.013088003732264042, t, erf_hp(0.00046277992078103125, t, erf_hm(0.0033736424520611763, t, 0.0015402256976813078)))))))) : F32)
1144
1145# [2.25, 2.75): t = a - 2.5
1146def erf_p3(+t: F32) -> F32:
1147 (0.9995930194854736 + erf_hp(2.846470792405853e-08, t, erf_hp(0.002178284339606762, t, erf_hm(0.005445694085210562, t, erf_hp(0.008350063115358353, t, erf_hm(0.008623690344393253, t, erf_hp(0.00611893879249692, t, erf_hm(0.002765072975307703, t, 0.0005077228415757418))))))) : F32)
1148
1149# [2.75, 3.25): t = a - 3.0
1150def erf_p4(+t: F32) -> F32:
1151 (0.9999778866767883 + erf_hp(2.2842012725732275e-08, t, erf_hp(0.00013925303937867284, t, erf_hm(0.0004177672672085464, t, erf_hp(0.0007891026325523853, t, erf_hm(0.0010437482269480824, t, erf_hp(0.001016448251903057, t, erf_hm(0.0007547215791419148, t, 0.0003925197524949908))))))) : F32)
1152
1153# [3.25, 3.92): t = a - 3.5
1154def erf_p5(+t: F32) -> F32:
1155 (0.9999992847442627 + erf_hm(2.7855115547481546e-08, t, erf_hp(5.400151167123113e-06, t, erf_hm(1.8894650565925986e-05, t, erf_hp(4.221926792524755e-05, t, erf_hm(6.778862734790891e-05, t, erf_hp(8.422814426012337e-05, t, erf_hm(7.95367595856078e-05, t, 4.151461325818673e-05))))))) : F32)
1156
1157def erf_sign(+x: F32, +e: F32) -> F32:
1158 pick(F32.is_lt(x, 0.0), F32.neg(e), e)
1159
1160# the pieces are chosen from the top down, so that only one polynomial is computed (each test is a match on a Bool
1161# parameter); b is "a is at least the piece's lower end"
1162def erf_i0(b: Bool, +x: F32, +a: F32) -> F32:
1163 match b:
1164 case True{}: erf_sign(x, erf_p0((a - 1.0 : F32)))
1165 case False{}: erf_small(x)
1166
1167def erf_i1(b: Bool, +x: F32, +a: F32) -> F32:
1168 match b:
1169 case True{}: erf_sign(x, erf_p1((a - 1.5 : F32)))
1170 case False{}: erf_i0(F32.is_ge(a, 0.84375), x, a)
1171
1172def erf_i2(b: Bool, +x: F32, +a: F32) -> F32:
1173 match b:
1174 case True{}: erf_sign(x, erf_p2((a - 2.0 : F32)))
1175 case False{}: erf_i1(F32.is_ge(a, 1.25), x, a)
1176
1177def erf_i3(b: Bool, +x: F32, +a: F32) -> F32:
1178 match b:
1179 case True{}: erf_sign(x, erf_p3((a - 2.5 : F32)))
1180 case False{}: erf_i2(F32.is_ge(a, 1.75), x, a)
1181
1182def erf_i4(b: Bool, +x: F32, +a: F32) -> F32:
1183 match b:
1184 case True{}: erf_sign(x, erf_p4((a - 3.0 : F32)))
1185 case False{}: erf_i3(F32.is_ge(a, 2.25), x, a)
1186
1187def erf_i5(b: Bool, +x: F32, +a: F32) -> F32:
1188 match b:
1189 case True{}: erf_sign(x, erf_p5((a - 3.5 : F32)))
1190 case False{}: erf_i4(F32.is_ge(a, 2.75), x, a)
1191
1192def erf_top(b: Bool, +x: F32, +a: F32) -> F32:
1193 match b:
1194 case True{}: erf_sign(x, 1.0)
1195 case False{}: erf_i5(F32.is_ge(a, 3.25), x, a)
1196
1197def erf(+x: F32) -> F32:
1198 +a = F32.abs(x)
1199 erf_top(F32.is_ge(a, 3.92), x, a)
1200
1201# GELU in its exact form, PyTorch's default (approximate="none"): gelu(x) = x/2 (1 + erf(x / sqrt 2)), in PyTorch's CPU
1202# order, (x * 0.5) * (1 + erf(x * 0.70710677)). Its gradient is cdf + x pdf, with cdf = 0.5 (1 + erf(x / sqrt 2)) and
1203# pdf = exp(-x^2 / 2) / sqrt(2 pi), computed as PyTorch does: 0.5 (1 + erf(x * 0.70710677)) and
1204# 0.3989423 * exp((x * x) * -0.5). The exponential is F32.exp, as in SiLU (so on a GPU the gradient would take CUDA's
1205# float exp).
1206def gelu_erf_cdf(+x: F32) -> F32:
1207 (0.5 * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32)
1208
1209def gelu_erf_v(+x: F32) -> F32:
1210 ((x * 0.5 : F32) * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32)
1211
1212def gelu_erf_d(+dy: F32, +x: F32) -> F32:
1213 +pdf = (0.3989422804014327 * F32.exp(((x * x : F32) * F32.neg(0.5) : F32)) : F32)
1214 (dy * (gelu_erf_cdf(x) + (x * pdf : F32) : F32) : F32)
1215
1216def GELUErf(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1217 ActIn(~gelu_erf_v, ~gelu_erf_d, n, c)
1218
1219# =====================================================================================================
1220# Dropout: each number is kept with probability 1 - p and scaled by 1 / (1 - p), or set to 0. The mask comes from a
1221# counter-based hash of (seed, step, index), so a run is reproducible and a twin can compute the same mask. The step
1222# counter is the layer's "parameter": the optimizer step only increments it, and checkpoints keep it.
1223# =====================================================================================================
1224
1225type DropP is Data:
1226 DropP{t: U32}
1227
1228# lowbias32 (Chris Wellons): an invertible mix of the 32 bits
1229def mix32(+x: U32) -> U32:
1230 +a = U32.xor(x, U32.shrn(x, 16n))
1231 +b = (a * 2146121005 : U32)
1232 +c = U32.xor(b, U32.shrn(b, 15n))
1233 +d = (c * 2221713035 : U32)
1234 U32.xor(d, U32.shrn(d, 16n))
1235
1236# the factor of number i: 1 / (1 - p) if its uniform u = (h >> 8) / 2^24 is >= p, else 0
1237def drop_k(+key: U32, +i: U32, +p: F32, +scale: F32) -> F32:
1238 pick(F32.is_lt((U32.to_f32(U32.shrn(mix32((key + i : U32)), 8n)) / 16777216.0 : F32), p), 0.0, scale)
1239
1240def drop_l(nleft: Nat, +i: U32, +key: U32, +p: F32, +scale: F32, a: Array<F32>) -> Array<F32>:
1241 match nleft:
1242 case 0n:
1243 a
1244 case 1n+q:
1245 drop_l(q, (i + 1 : U32), key, p, scale, Array.set(F32, a, i, drop_k(key, i, p, scale)))
1246
1247def drop_mask(+n: Nat, +c: Nat, +seed: U32, +t: U32, +p: F32) -> TA.Mat<n, c>:
1248 TA.Mat{drop_l(Nat.mul(n, c), 0, mix32(U32.xor(seed, mix32(t))), p, (1.0 / (1.0 - p : F32) : F32), zeros_a(Nat.mul(n, c)))}
1249
1250def drop_f2(-n: Nat, -c: Nat, +t: U32, r: TA.MPair<n, c>) -> Fwd<DropP, TA.Mat<n, c>, TA.Mat<n, c>>:
1251 match r:
1252 case TA.MPair{y, k}:
1253 Fwd{DropP{t}, y, k}
1254
1255def drop_fwd(+n: Nat, +c: Nat, +seed: U32, +p: F32, d: DropP, x: TA.Mat<n, c>) -> Fwd<DropP, TA.Mat<n, c>, TA.Mat<n, c>>:
1256 match d:
1257 case DropP{+t}:
1258 drop_f2(n, c, t, Ew.zip(~f32_mul, n, c, x, drop_mask(n, c, seed, t, p)))
1259
1260def drop_b2(-n: Nat, -c: Nat, d: DropP, g: DropP, r: TA.MPair<n, c>) -> Bwd<DropP, TA.Mat<n, c>>:
1261 match r:
1262 case TA.MPair{dx, k}:
1263 Bwd{d, dx, g}
1264
1265def drop_step(+op: Opt, d: DropP, g: DropP, st: OptState<DropP>) -> Step<DropP>:
1266 match d:
1267 case DropP{t}:
1268 Step{DropP{(t + 1 : U32)}, g, st}
1269
1270def drop_dump(d: DropP) -> PL<DropP>:
1271 match d:
1272 case DropP{+t}:
1273 PL{DropP{t}, [bitcast(t)]}
1274
1275def drop_load_fin(+off: U32, r: Array<F32> & F32) -> PA<DropP>:
1276 match r:
1277 case Tuple{a, v}: PA{DropP{bits(v)}, a, (off + 1 : U32)}
1278
1279def drop_load(d: DropP, xs: Array<F32>, +off: U32) -> PA<DropP>:
1280 drop_load_fin(off, Array.get(F32, xs, off))
1281
1282# the counter, bitcast to one F32 as dump writes it
1283def drop_store(d: DropP, xs: Array<F32>, +off: U32) -> PA<DropP>:
1284 match d:
1285 case DropP{+c}: PA{DropP{c}, Array.set(F32, xs, off, bitcast(c)), (off + 1 : U32)}
1286
1287def Dropout(+n: Nat, +c: Nat, +p: F32, +seed: U32) -> Layer<DropP, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1288 Layer{d => x => drop_fwd(n, c, seed, p, d, x), d => s => dy => g => drop_b2(n, c, d, g, Ew.zip(~f32_mul, n, c, dy, s)), op => d => g => st => drop_step(op, d, g, st), u => DropP{0}, d => drop_dump(d), d => xs => off => drop_load(d, xs, off), d => xs => off => drop_store(d, xs, off), u => 1}
1289
1290# =====================================================================================================
1291# Row kernels over arrays: loops over a row of a flat array by index. Every sum runs left to right from +0.0 and every
1292# expression keeps its order of operations, the ones the row-wise layers had when they worked on lists of rows, so
1293# they give the same bits (NOTES.md, M6)
1294# =====================================================================================================
1295
1296type D2 is Type:
1297 D2{a: Array<F32>, b: Array<F32>, s: F32}
1298
1299type A2 is Type:
1300 A2{a: Array<F32>, b: Array<F32>}
1301
1302# a[ia]·b[ib] + a[ia + sa]·b[ib + sb] + ... (k terms) added to acc left to right; ra and rb hold the
1303# first pair, already read
1304def sdot(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> D2:
1305 match k ra rb:
1306 case 0n Tuple{a, x} Tuple{b, y}:
1307 D2{a, b, acc}
1308 case 1n+p Tuple{a, +x} Tuple{b, +y}:
1309 sdot(p, (ia + sa : U32), (ib + sb : U32), sa, sb, (acc + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32)))
1310
1311def dot0(+k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a: Array<F32>, b: Array<F32>) -> D2:
1312 sdot(k, ia, ib, sa, sb, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib))
1313
1314
1315type AS is Type:
1316 AS{a: Array<F32>, s: F32}
1317
1318# the largest of a[idx] (held in r) and the k numbers after it: max(max(a0, a1), a2) ..., from the first
1319def pmax(k: Nat, +idx: U32, cur: F32, r: Array<F32> & F32) -> AS:
1320 match k r:
1321 case 0n Tuple{a, +x}:
1322 AS{a, F32.max(cur, x)}
1323 case 1n+q Tuple{a, +x}:
1324 pmax(q, (idx + 1 : U32), F32.max(cur, x), Array.get(F32, a, (idx + 1 : U32)))
1325
1326def pmax0(k: Nat, +idx: U32, r: Array<F32> & F32) -> AS:
1327 match k r:
1328 case 0n Tuple{a, +x}:
1329 AS{a, x}
1330 case 1n+q Tuple{a, +x}:
1331 pmax(q, (idx + 1 : U32), x, Array.get(F32, a, (idx + 1 : U32)))
1332
1333# a[j] = exp(a[j] - m) for j = idx..idx + k, their sum added to acc in order
1334def pexp(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array<F32> & F32) -> AS:
1335 match k r:
1336 case 0n Tuple{a, +x}:
1337 +e = F32.exp((x - m : F32))
1338 AS{Array.set(F32, a, idx, e), (acc + e : F32)}
1339 case 1n+q Tuple{a, +x}:
1340 +e = F32.exp((x - m : F32))
1341 pexp(q, (idx + 1 : U32), m, (acc + e : F32), Array.get(F32, Array.set(F32, a, idx, e), (idx + 1 : U32)))
1342
1343# the array and how many numbers the last loop of a softmax divided
1344type PC is Type:
1345 PC{a: Array<F32>, c: Nat}
1346
1347def pc_a(r: PC) -> Array<F32>:
1348 match r:
1349 case PC{a, c}:
1350 a
1351
1352# a[j] = a[j] / s for j = idx..idx + k, counted onto c
1353def pdiv(k: Nat, +idx: U32, +s: F32, c: Nat, r: Array<F32> & F32) -> PC:
1354 match k r:
1355 case 0n Tuple{a, +x}:
1356 PC{Array.set(F32, a, idx, (x / s : F32)), 1n+c}
1357 case 1n+q Tuple{a, +x}:
1358 pdiv(q, (idx + 1 : U32), s, 1n+c, Array.get(F32, Array.set(F32, a, idx, (x / s : F32)), (idx + 1 : U32)))
1359
1360def psm3(k: Nat, +off: U32, r: AS) -> PC:
1361 match r:
1362 case AS{a, +s}:
1363 pdiv(k, off, s, 0n, Array.get(F32, a, off))
1364
1365def psm2(+k: Nat, +off: U32, r: AS) -> PC:
1366 match r:
1367 case AS{a, +m}:
1368 psm3(k, off, pexp(k, off, m, 0.0, Array.get(F32, a, off)))
1369
1370# softmax in place over the n numbers of p from off (n >= 1), as PyTorch: exp(x - max) / sum; with the count of the
1371# numbers it divided (the law causal reads it)
1372def psmc(+n: Nat, +off: U32, p: Array<F32>) -> PC:
1373 psm2(Nat.sub(n, 1n), off, pmax0(Nat.sub(n, 1n), off, Array.get(F32, p, off)))
1374
1375def psm(+n: Nat, +off: U32, p: Array<F32>) -> Array<F32>:
1376 pc_a(psmc(n, off, p))
1377
1378# dst[di..] = src[si..], cnt + 1 numbers; r holds src[si], already read
1379def copy_n(cnt: Nat, +si: U32, +di: U32, o: Array<F32>, r: Array<F32> & F32) -> A2:
1380 match cnt r:
1381 case 0n Tuple{s, +x}:
1382 A2{s, Array.set(F32, o, di, x)}
1383 case 1n+n Tuple{s, +x}:
1384 copy_n(n, (si + 1 : U32), (di + 1 : U32), Array.set(F32, o, di, x), Array.get(F32, s, (si + 1 : U32)))
1385
1386# =====================================================================================================
1387# Row-wise layers: softmax, log-softmax, RMSNorm, LayerNorm (each row of an n x c matrix on its own)
1388# =====================================================================================================
1389
1390# rows i.. of x (c wide, left more after i), each softmaxed in place (psm)
1391def smf_rows(left: Nat, +i: U32, +c: Nat, x: Array<F32>) -> Array<F32>:
1392 match left:
1393 case 0n:
1394 psm(c, (i * U32.from_nat(c) : U32), x)
1395 case 1n+q:
1396 smf_rows(q, (i + 1 : U32), c, psm(c, (i * U32.from_nat(c) : U32), x))
1397
1398# the output, and a copy of it saved for the backward pass
1399def out_saved(-n: Nat, -c: Nat, r: TA.MM2<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1400 match r:
1401 case TA.MM2{a, b}:
1402 Fwd{Unit{}, a, b}
1403
1404def sm_fwd(+n: Nat, +c: Nat, x: TA.Mat<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1405 match x:
1406 case TA.Mat{a}:
1407 out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{smf_rows(Nat.sub(n, 1n), 0, c, a)}))
1408
1409# dy[j] = y[j] * (dy[j] - s), j = idx..idx + k; rd and ry hold dy[idx] and y[idx]
1410def smb_fix(k: Nat, +idx: U32, +s: F32, rd: Array<F32> & F32, ry: Array<F32> & F32) -> A2:
1411 match k rd ry:
1412 case 0n Tuple{d, +g} Tuple{y, +v}:
1413 A2{Array.set(F32, d, idx, (v * (g - s : F32) : F32)), y}
1414 case 1n+q Tuple{d, +g} Tuple{y, +v}:
1415 smb_fix(q, (idx + 1 : U32), s, Array.get(F32, Array.set(F32, d, idx, (v * (g - s : F32) : F32)), (idx + 1 : U32)), Array.get(F32, y, (idx + 1 : U32)))
1416
1417def smb_row(+c: Nat, +off: U32, r: D2) -> A2:
1418 match r:
1419 case D2{d, y, +s}:
1420 smb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off))
1421
1422# softmax backward, as PyTorch's: dx = y * (dy - sum(dy * y)), row after row, in place over dy
1423def smb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2:
1424 match left st:
1425 case 0n A2{d, y}:
1426 smb_row(c, (i * U32.from_nat(c) : U32), dot0(c, (i * U32.from_nat(c) : U32), (i * U32.from_nat(c) : U32), 1, 1, d, y))
1427 case 1n+q A2{d, y}:
1428 smb_rows(q, (i + 1 : U32), c, smb_row(c, (i * U32.from_nat(c) : U32), dot0(c, (i * U32.from_nat(c) : U32), (i * U32.from_nat(c) : U32), 1, 1, d, y)))
1429
1430def dx_fin(-n: Nat, -c: Nat, r: A2) -> Bwd<Unit, TA.Mat<n, c>>:
1431 match r:
1432 case A2{d, y}:
1433 Bwd{Unit{}, TA.Mat{d}, Unit{}}
1434
1435def sm_bwd(+n: Nat, +c: Nat, s: TA.Mat<n, c>, dy: TA.Mat<n, c>) -> Bwd<Unit, TA.Mat<n, c>>:
1436 match s dy:
1437 case TA.Mat{y} TA.Mat{d}:
1438 dx_fin(n, c, smb_rows(Nat.sub(n, 1n), 0, c, A2{d, y}))
1439
1440# softmax over each row; saves its output
1441def Softmax(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1442 Layer{p => x => sm_fwd(n, c, x), p => s => dy => g => sm_bwd(n, c, s, dy), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
1443
1444# log_softmax(x) = (x - max) - log(sum(exp(x - max))), as PyTorch's CPU kernel
1445
1446# acc + exp(a[j] - m) for j = idx..idx + k, in order
1447def lsum(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array<F32> & F32) -> AS:
1448 match k r:
1449 case 0n Tuple{a, +x}:
1450 AS{a, (acc + F32.exp((x - m : F32)) : F32)}
1451 case 1n+q Tuple{a, +x}:
1452 lsum(q, (idx + 1 : U32), m, (acc + F32.exp((x - m : F32)) : F32), Array.get(F32, a, (idx + 1 : U32)))
1453
1454# a[j] = (a[j] - m) - l for j = idx..idx + k
1455def lw(k: Nat, +idx: U32, +m: F32, +l: F32, r: Array<F32> & F32) -> Array<F32>:
1456 match k r:
1457 case 0n Tuple{a, +x}:
1458 Array.set(F32, a, idx, ((x - m : F32) - l : F32))
1459 case 1n+q Tuple{a, +x}:
1460 lw(q, (idx + 1 : U32), m, l, Array.get(F32, Array.set(F32, a, idx, ((x - m : F32) - l : F32)), (idx + 1 : U32)))
1461
1462def lsm_r3(+c: Nat, +off: U32, +m: F32, r: AS) -> Array<F32>:
1463 match r:
1464 case AS{a, +s}:
1465 lw(Nat.sub(c, 1n), off, m, F32.log(s), Array.get(F32, a, off))
1466
1467def lsm_r2(+c: Nat, +off: U32, r: AS) -> Array<F32>:
1468 match r:
1469 case AS{a, +m}:
1470 lsm_r3(c, off, m, lsum(Nat.sub(c, 1n), off, m, 0.0, Array.get(F32, a, off)))
1471
1472def lsm_row(+c: Nat, +off: U32, x: Array<F32>) -> Array<F32>:
1473 lsm_r2(c, off, pmax0(Nat.sub(c, 1n), off, Array.get(F32, x, off)))
1474
1475def lsmf_rows(left: Nat, +i: U32, +c: Nat, x: Array<F32>) -> Array<F32>:
1476 match left:
1477 case 0n:
1478 lsm_row(c, (i * U32.from_nat(c) : U32), x)
1479 case 1n+q:
1480 lsmf_rows(q, (i + 1 : U32), c, lsm_row(c, (i * U32.from_nat(c) : U32), x))
1481
1482def lsm_fwd(+n: Nat, +c: Nat, x: TA.Mat<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1483 match x:
1484 case TA.Mat{a}:
1485 out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{lsmf_rows(Nat.sub(n, 1n), 0, c, a)}))
1486
1487# acc + a[j] for j = idx..idx + k, in order
1488def asum(k: Nat, +idx: U32, acc: F32, r: Array<F32> & F32) -> AS:
1489 match k r:
1490 case 0n Tuple{a, +x}:
1491 AS{a, (acc + x : F32)}
1492 case 1n+q Tuple{a, +x}:
1493 asum(q, (idx + 1 : U32), (acc + x : F32), Array.get(F32, a, (idx + 1 : U32)))
1494
1495# dy[j] = dy[j] - exp(y[j]) * s, j = idx..idx + k
1496def lsmb_fix(k: Nat, +idx: U32, +s: F32, rd: Array<F32> & F32, ry: Array<F32> & F32) -> A2:
1497 match k rd ry:
1498 case 0n Tuple{d, +g} Tuple{y, +v}:
1499 A2{Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), y}
1500 case 1n+q Tuple{d, +g} Tuple{y, +v}:
1501 lsmb_fix(q, (idx + 1 : U32), s, Array.get(F32, Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), (idx + 1 : U32)), Array.get(F32, y, (idx + 1 : U32)))
1502
1503def lsmb_row(+c: Nat, +off: U32, y: Array<F32>, r: AS) -> A2:
1504 match r:
1505 case AS{d, +s}:
1506 lsmb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off))
1507
1508# dx = dy - exp(y) * sum(dy), row after row, in place over dy
1509def lsmb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2:
1510 match left st:
1511 case 0n A2{d, y}:
1512 lsmb_row(c, (i * U32.from_nat(c) : U32), y, asum(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), 0.0, Array.get(F32, d, (i * U32.from_nat(c) : U32))))
1513 case 1n+q A2{d, y}:
1514 lsmb_rows(q, (i + 1 : U32), c, lsmb_row(c, (i * U32.from_nat(c) : U32), y, asum(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), 0.0, Array.get(F32, d, (i * U32.from_nat(c) : U32)))))
1515
1516def lsm_bwd(+n: Nat, +c: Nat, s: TA.Mat<n, c>, dy: TA.Mat<n, c>) -> Bwd<Unit, TA.Mat<n, c>>:
1517 match s dy:
1518 case TA.Mat{y} TA.Mat{d}:
1519 dx_fin(n, c, lsmb_rows(Nat.sub(n, 1n), 0, c, A2{d, y}))
1520
1521def LogSoftmax(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1522 Layer{p => x => lsm_fwd(n, c, x), p => s => dy => g => lsm_bwd(n, c, s, dy), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
1523
1524# ---- RMSNorm: y = g * (x / sqrt(mean(x^2) + eps)); its parameter is the gain g (1 x d) ----
1525
1526# one optimizer step on a layer whose parameter is a single matrix
1527def mat_step_fin(-r: Nat, -c: Nat, q: MatStep<r, c>) -> Step<TA.Mat<r, c>>:
1528 match q:
1529 case MatStep{w, g, st}: Step{w, g, st}
1530
1531def mat_step(+r: Nat, +c: Nat, +op: Opt, p: TA.Mat<r, c>, g: TA.Mat<r, c>, st: OptState<TA.Mat<r, c>>) -> Step<TA.Mat<r, c>>:
1532 mat_step_fin(r, c, Opt.mat(r, c, op, p, g, st))
1533
1534def g_add_fin(-r: Nat, -c: Nat, r2: TA.MPair<r, c>) -> TA.Mat<r, c>:
1535 match r2:
1536 case TA.MPair{a, b}:
1537 a
1538
1539# a + b, dropping b
1540def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
1541 g_add_fin(r, c, Mat.add(r, c, a, b))
1542
1543def dup_fin(-r: Nat, -c: Nat, m: TA.MM2<r, c>) -> Two<TA.Mat<r, c>, TA.Mat<r, c>>:
1544 match m:
1545 case TA.MM2{a, b}:
1546 Two{a, b}
1547
1548# two copies of a matrix
1549def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> Two<TA.Mat<r, c>, TA.Mat<r, c>>:
1550 dup_fin(r, c, TA.Mat.clone(r, c, m))
1551
1552def ml_list(-r: Nat, -c: Nat, ml: TA.ML<r, c>) -> List<&2, F32>:
1553 match ml:
1554 case TA.ML{m, xs}:
1555 xs
1556
1557# ---- RMSNorm layer, a row at a time over the arrays: x (n x d), the gain g (1 x d), the output y. Each sum is added
1558# left to right from +0.0, each product in the order the list version multiplied, so the bits are the list version's.
1559# ----
1560
1561# acc + a[i]·a[i] + ... over k + 1 numbers from i, in order; r holds a[i], already read
1562def ssq(k: Nat, +i: U32, acc: F32, r: Array<F32> & F32) -> AS:
1563 match k r:
1564 case 0n Tuple{a, +x}:
1565 AS{a, (acc + (x * x : F32) : F32)}
1566 case 1n+p Tuple{a, +x}:
1567 ssq(p, (i + 1 : U32), (acc + (x * x : F32) : F32), Array.get(F32, a, (i + 1 : U32)))
1568
1569# 1 / rms of a row whose sum of squares is s
1570def rms_inv(+s: F32, +df: F32, +eps: F32) -> F32:
1571 (1.0 / F32.sqrt(((s / df : F32) + eps : F32)) : F32)
1572
1573type R3 is Type:
1574 R3{x: Array<F32>, g: Array<F32>, y: Array<F32>}
1575
1576# y[off + j] = g[j] * (x[off + j] * r), j = jj..jj + k; rx and rg hold x[off + jj] and g[jj]
1577def rms_w(k: Nat, +j: U32, +off: U32, +r: F32, y: Array<F32>, rx: Array<F32> & F32, rg: Array<F32> & F32) -> R3:
1578 match k rx rg:
1579 case 0n Tuple{x, +a} Tuple{g, +b}:
1580 R3{x, g, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32))}
1581 case 1n+p Tuple{x, +a} Tuple{g, +b}:
1582 rms_w(p, (j + 1 : U32), off, r, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
1583
1584def rmsf_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, g: Array<F32>, y: Array<F32>, r: AS) -> R3:
1585 match r:
1586 case AS{x, +s}:
1587 rms_w(Nat.sub(d, 1n), 0, off, rms_inv(s, df, eps), y, Array.get(F32, x, off), Array.get(F32, g, 0))
1588
1589def rmsf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R3) -> R3:
1590 match s:
1591 case R3{x, g, y}:
1592 rmsf_row2(d, off, df, eps, g, y, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off)))
1593
1594# rows i.. (left + 1 of them, d wide)
1595def rmsf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R3) -> R3:
1596 match left:
1597 case 0n:
1598 rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1599 case 1n+p:
1600 rmsf_rows(p, (i + 1 : U32), d, df, eps, rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1601
1602def rms_f2(-n: Nat, -d: Nat, s: R3) -> Fwd<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1603 match s:
1604 case R3{x, g, y}:
1605 Fwd{TA.Mat{g}, TA.Mat{y}, TA.Mat{x}}
1606
1607def rms_fwd(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat<n, d>) -> Fwd<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1608 match p x:
1609 case TA.Mat{g} TA.Mat{xa}:
1610 rms_f2(n, d, rmsf_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R3{xa, g, zeros_a(Nat.mul(n, d))}))
1611
1612# with r = 1/rms, xh = x*r, dxh = dy*g: dx = r * (dxh - xh * mean(dxh * xh)); dg += dy * xh, row after row
1613type R5 is Type:
1614 R5{x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, dg: Array<F32>}
1615
1616type R3S is Type:
1617 R3S{x: Array<F32>, dy: Array<F32>, g: Array<F32>, s: F32}
1618
1619# acc + dxh[j]·xh[j] + ..., in order; rx, rd, rg hold x[off + jj], dy[off + jj], g[jj]
1620def rmsb_dot(k: Nat, +j: U32, +off: U32, +r: F32, acc: F32, rx: Array<F32> & F32, rd: Array<F32> & F32, rg: Array<F32> & F32) -> R3S:
1621 match k rx rd rg:
1622 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1623 R3S{x, dy, g, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32)}
1624 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1625 rmsb_dot(p, (j + 1 : U32), off, r, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
1626
1627def dg_add(+j: U32, +v: F32, r: Array<F32> & F32) -> Array<F32>:
1628 match r:
1629 case Tuple{dg, +old}:
1630 Array.set(F32, dg, j, (old + v : F32))
1631
1632# dx[off + j] = ((dy[j]·g[j]) - (x[j]·r)·m)·r and dg[j] += dy[j]·(x[j]·r), j = jj..jj + k
1633def rmsb_w(k: Nat, +j: U32, +off: U32, +r: F32, +m: F32, dx: Array<F32>, dg: Array<F32>, rx: Array<F32> & F32, rd: Array<F32> & F32, rg: Array<F32> & F32) -> R5:
1634 match k rx rd rg:
1635 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1636 +xh = (a * r : F32)
1637 R5{x, dy, g, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), dg_add(j, (b * xh : F32), Array.get(F32, dg, j))}
1638 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1639 +xh = (a * r : F32)
1640 rmsb_w(p, (j + 1 : U32), off, r, m, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), dg_add(j, (b * xh : F32), Array.get(F32, dg, j)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
1641
1642def rmsb_row3(+d: Nat, +off: U32, +df: F32, +r: F32, dx: Array<F32>, dg: Array<F32>, q: R3S) -> R5:
1643 match q:
1644 case R3S{x, dy, g, +s}:
1645 rmsb_w(Nat.sub(d, 1n), 0, off, r, (s / df : F32), dx, dg, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0))
1646
1647def rmsb_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, dg: Array<F32>, q: AS) -> R5:
1648 match q:
1649 case AS{x, +s}:
1650 +r = rms_inv(s, df, eps)
1651 rmsb_row3(d, off, df, r, dx, dg, rmsb_dot(Nat.sub(d, 1n), 0, off, r, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)))
1652
1653def rmsb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R5) -> R5:
1654 match s:
1655 case R5{x, dy, g, dx, dg}:
1656 rmsb_row2(d, off, df, eps, dy, g, dx, dg, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off)))
1657
1658def rmsb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R5) -> R5:
1659 match left:
1660 case 0n:
1661 rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1662 case 1n+p:
1663 rmsb_rows(p, (i + 1 : U32), d, df, eps, rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1664
1665def gain_bwd_fin(-n: Nat, -d: Nat, -P: Type, p: P, dx: TA.Mat<n, d>, g: P) -> Bwd<P, TA.Mat<n, d>>:
1666 Bwd{p, dx, g}
1667
1668def rms_b2(+d: Nat, -n: Nat, gacc: TA.Mat<1n, d>, s: R5) -> Bwd<TA.Mat<1n, d>, TA.Mat<n, d>>:
1669 match s:
1670 case R5{x, dy, g, dx, dg}:
1671 gain_bwd_fin(n, d, TA.Mat<1n, d>, TA.Mat{g}, TA.Mat{dx}, Mat.sum2(1n, d, gacc, TA.Mat{dg}))
1672
1673def rms_bwd(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat<n, d>, dy: TA.Mat<n, d>, gacc: TA.Mat<1n, d>) -> Bwd<TA.Mat<1n, d>, TA.Mat<n, d>>:
1674 match p x dy:
1675 case TA.Mat{g} TA.Mat{xa} TA.Mat{da}:
1676 rms_b2(d, n, gacc, rmsb_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R5{xa, da, g, zeros_a(Nat.mul(n, d)), zeros_a(d)}))
1677
1678# RMSNorm over rows of d numbers, with a learned gain (Llama's norm); saves its input
1679def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1680 Layer{p => x => rms_fwd(n, d, eps, p, x), p => s => dy => g => rms_bwd(n, d, eps, p, s, dy, g), op => p => g => st => mat_step(1n, d, op, p, g, st), u => TA.Mat.zeros(1n, d), p => Mat.dump(1n, d, p), p => xs => off => Mat.restore(1n, d, p, xs, off), p => xs => off => Mat.store(1n, d, p, xs, off), u => U32.from_nat(d)}
1681
1682# ---- LayerNorm: y = (x - mean) / sqrt(var + eps) * g + b, var the biased variance ----
1683
1684type LN<-d: Nat> is Type:
1685 LN{g: TA.Mat<1n, d>, b: TA.Mat<1n, d>}
1686
1687# acc + (a[j] - mu)·(a[j] - mu) for j = idx..idx + k, in order
1688def csq(k: Nat, +idx: U32, +mu: F32, acc: F32, r: Array<F32> & F32) -> AS:
1689 match k r:
1690 case 0n Tuple{a, +x}:
1691 AS{a, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32)}
1692 case 1n+q Tuple{a, +x}:
1693 csq(q, (idx + 1 : U32), mu, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32), Array.get(F32, a, (idx + 1 : U32)))
1694
1695type ASM is Type:
1696 ASM{a: Array<F32>, mu: F32, r: F32}
1697
1698def ln_mr2(+df: F32, +eps: F32, +mu: F32, q: AS) -> ASM:
1699 match q:
1700 case AS{a, +v}:
1701 ASM{a, mu, (1.0 / F32.sqrt(((v / df : F32) + eps : F32)) : F32)}
1702
1703def ln_mr1(+d: Nat, +off: U32, +df: F32, +eps: F32, q: AS) -> ASM:
1704 match q:
1705 case AS{a, +s}:
1706 ln_mr2(df, eps, (s / df : F32), csq(Nat.sub(d, 1n), off, (s / df : F32), 0.0, Array.get(F32, a, off)))
1707
1708# a row's mean and 1 / sqrt(var + eps), as ln_xh and ln_r computed them
1709def ln_mr(+d: Nat, +off: U32, +df: F32, +eps: F32, x: Array<F32>) -> ASM:
1710 ln_mr1(d, off, df, eps, asum(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off)))
1711
1712type R4 is Type:
1713 R4{x: Array<F32>, g: Array<F32>, b: Array<F32>, y: Array<F32>}
1714
1715# y[off + j] = ((x - mu) * r) * g[j] + b[j], j = jj..jj + k
1716def lnf_w(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, y: Array<F32>, rx: Array<F32> & F32, rg: Array<F32> & F32, rb: Array<F32> & F32) -> R4:
1717 match k rx rg rb:
1718 case 0n Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}:
1719 R4{x, g, b, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32))}
1720 case 1n+q Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}:
1721 lnf_w(q, (j + 1 : U32), off, mu, r, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)), Array.get(F32, b, (j + 1 : U32)))
1722
1723def lnf_row2(+d: Nat, +off: U32, g: Array<F32>, b: Array<F32>, y: Array<F32>, m: ASM) -> R4:
1724 match m:
1725 case ASM{x, +mu, +r}:
1726 lnf_w(Nat.sub(d, 1n), 0, off, mu, r, y, Array.get(F32, x, off), Array.get(F32, g, 0), Array.get(F32, b, 0))
1727
1728def lnf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4) -> R4:
1729 match s:
1730 case R4{x, g, b, y}:
1731 lnf_row2(d, off, g, b, y, ln_mr(d, off, df, eps, x))
1732
1733def lnf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4) -> R4:
1734 match left:
1735 case 0n:
1736 lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1737 case 1n+q:
1738 lnf_rows(q, (i + 1 : U32), d, df, eps, lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1739
1740def ln_f2(-n: Nat, -d: Nat, s: R4) -> Fwd<LN<d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1741 match s:
1742 case R4{x, g, b, y}:
1743 Fwd{LN{TA.Mat{g}, TA.Mat{b}}, TA.Mat{y}, TA.Mat{x}}
1744
1745def ln_fwd(+n: Nat, +d: Nat, +eps: F32, p: LN<d>, x: TA.Mat<n, d>) -> Fwd<LN<d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1746 match p x:
1747 case LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{xa}:
1748 ln_f2(n, d, lnf_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R4{xa, g, b, zeros_a(Nat.mul(n, d))}))
1749
1750# with xh = (x - mu) * r and dxh = dy * g: dx = r * ((dxh - mean(dxh)) - xh * mean(dxh * xh)); dg += dy * xh; db += dy
1751type R3SS is Type:
1752 R3SS{x: Array<F32>, dy: Array<F32>, g: Array<F32>, sa: F32, sb: F32}
1753
1754# sa + dxh[j] and sb + dxh[j]·xh[j], j = jj..jj + k, each in order
1755def lnb_sums(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, sa: F32, sb: F32, rx: Array<F32> & F32, rd: Array<F32> & F32, rg: Array<F32> & F32) -> R3SS:
1756 match k rx rd rg:
1757 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1758 R3SS{x, dy, g, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32)}
1759 case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1760 lnb_sums(q, (j + 1 : U32), off, mu, r, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
1761
1762def ln_dw(+j: U32, +v: F32, +w: F32, st: A2) -> A2:
1763 match st:
1764 case A2{dg, db}:
1765 A2{dg_add(j, v, Array.get(F32, dg, j)), dg_add(j, w, Array.get(F32, db, j))}
1766
1767type R4B is Type:
1768 R4B{x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, w: A2}
1769
1770# dx[off + j] = (((dxh - ma) - xh * mb) * r), dg[j] += dy * xh, db[j] += dy, j = jj..jj + k (ma, mb the two means)
1771def lnb_w(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, +ma: F32, +mb: F32, dx: Array<F32>, w: A2, rx: Array<F32> & F32, rd: Array<F32> & F32, rg: Array<F32> & F32) -> R4B:
1772 match k rx rd rg:
1773 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1774 +xh = ((a - mu : F32) * r : F32)
1775 R4B{x, dy, g, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), ln_dw(j, (b * xh : F32), b, w)}
1776 case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1777 +xh = ((a - mu : F32) * r : F32)
1778 lnb_w(q, (j + 1 : U32), off, mu, r, ma, mb, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), ln_dw(j, (b * xh : F32), b, w), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)))
1779
1780def lnb_row3(+d: Nat, +off: U32, +df: F32, +mu: F32, +r: F32, dx: Array<F32>, w: A2, q: R3SS) -> R4B:
1781 match q:
1782 case R3SS{x, dy, g, +sa, +sb}:
1783 lnb_w(Nat.sub(d, 1n), 0, off, mu, r, (sa / df : F32), (sb / df : F32), dx, w, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0))
1784
1785def lnb_row2(+d: Nat, +off: U32, +df: F32, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, w: A2, m: ASM) -> R4B:
1786 match m:
1787 case ASM{x, +mu, +r}:
1788 lnb_row3(d, off, df, mu, r, dx, w, lnb_sums(Nat.sub(d, 1n), 0, off, mu, r, 0.0, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)))
1789
1790def lnb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4B) -> R4B:
1791 match s:
1792 case R4B{x, dy, g, dx, w}:
1793 lnb_row2(d, off, df, dy, g, dx, w, ln_mr(d, off, df, eps, x))
1794
1795def lnb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4B) -> R4B:
1796 match left:
1797 case 0n:
1798 lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1799 case 1n+q:
1800 lnb_rows(q, (i + 1 : U32), d, df, eps, lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1801
1802def ln_b2(+d: Nat, -n: Nat, b: TA.Mat<1n, d>, gg: LN<d>, s: R4B) -> Bwd<LN<d>, TA.Mat<n, d>>:
1803 match gg s:
1804 case LN{ga, gb} R4B{x, dy, g, dx, A2{dg, db}}:
1805 gain_bwd_fin(n, d, LN<d>, LN{TA.Mat{g}, b}, TA.Mat{dx}, LN{Mat.sum2(1n, d, ga, TA.Mat{dg}), Mat.sum2(1n, d, gb, TA.Mat{db})})
1806
1807def ln_bwd(+n: Nat, +d: Nat, +eps: F32, p: LN<d>, x: TA.Mat<n, d>, dy: TA.Mat<n, d>, gg: LN<d>) -> Bwd<LN<d>, TA.Mat<n, d>>:
1808 match p x dy:
1809 case LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da}:
1810 ln_b2(d, n, b, gg, lnb_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R4B{xa, da, g, zeros_a(Nat.mul(n, d)), A2{zeros_a(d), zeros_a(d)}}))
1811
1812def ln_pair(-d: Nat, p: LN<d>) -> Two<TA.Mat<1n, d>, TA.Mat<1n, d>>:
1813 match p:
1814 case LN{w, b}: Two{w, b}
1815
1816def ln_unpair(-d: Nat, p: Two<TA.Mat<1n, d>, TA.Mat<1n, d>>) -> LN<d>:
1817 match p:
1818 case Two{w, b}: LN{w, b}
1819
1820def ln_s2(-d: Nat, a: MatStep<1n, d>, c: MatStep<1n, d>) -> Step<LN<d>>:
1821 match a c:
1822 case MatStep{w, gw, sw} MatStep{b, gb, sb}:
1823 Step{LN{w, b}, LN{gw, gb}, OptState.map(Two<TA.Mat<1n, d>, TA.Mat<1n, d>>, LN<d>, OptState.join(TA.Mat<1n, d>, TA.Mat<1n, d>, sw, sb), x => ln_unpair(d, x), x => ln_unpair(d, x))}
1824
1825def ln_step2(+d: Nat, +op: Opt, p: LN<d>, g: LN<d>, st: Two<OptState<TA.Mat<1n, d>>, OptState<TA.Mat<1n, d>>>) -> Step<LN<d>>:
1826 match p g st:
1827 case LN{w, b} LN{gw, gb} Two{sw, sb}:
1828 ln_s2(d, Opt.mat(1n, d, op, w, gw, sw), Opt.mat(1n, d, op, b, gb, sb))
1829
1830def ln_step(+d: Nat, +op: Opt, p: LN<d>, g: LN<d>, st: OptState<LN<d>>) -> Step<LN<d>>:
1831 ln_step2(d, op, p, g, OptState.split(TA.Mat<1n, d>, TA.Mat<1n, d>, OptState.map(LN<d>, Two<TA.Mat<1n, d>, TA.Mat<1n, d>>, st, x => ln_pair(d, x), x => ln_pair(d, x))))
1832
1833def ln_dump2(-d: Nat, r: PL<Two<TA.Mat<1n, d>, TA.Mat<1n, d>>>) -> PL<LN<d>>:
1834 match r:
1835 case PL{Two{g, b}, xs}:
1836 PL{LN{g, b}, xs}
1837
1838def ln_dump(+d: Nat, p: LN<d>) -> PL<LN<d>>:
1839 match p:
1840 case LN{g, b}:
1841 ln_dump2(d, pl_pair(TA.Mat<1n, d>, TA.Mat<1n, d>, Mat.dump(1n, d, g), Mat.dump(1n, d, b)))
1842
1843def ln_load3(-d: Nat, r: PA<Two<TA.Mat<1n, d>, TA.Mat<1n, d>>>) -> PA<LN<d>>:
1844 match r:
1845 case PA{Two{g, b}, xs, off}: PA{LN{g, b}, xs, off}
1846
1847def ln_load2(+d: Nat, b: TA.Mat<1n, d>, r: PA<TA.Mat<1n, d>>) -> PA<LN<d>>:
1848 match r:
1849 case PA{g, xs, off}:
1850 ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.restore(1n, d, b, xs, off)))
1851
1852def ln_load(+d: Nat, p: LN<d>, xs: Array<F32>, off: U32) -> PA<LN<d>>:
1853 match p:
1854 case LN{g, b}: ln_load2(d, b, Mat.restore(1n, d, g, xs, off))
1855
1856def ln_store2(+d: Nat, b: TA.Mat<1n, d>, r: PA<TA.Mat<1n, d>>) -> PA<LN<d>>:
1857 match r:
1858 case PA{g, xs, off}:
1859 ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.store(1n, d, b, xs, off)))
1860
1861def ln_store(+d: Nat, p: LN<d>, xs: Array<F32>, off: U32) -> PA<LN<d>>:
1862 match p:
1863 case LN{g, b}: ln_store2(d, b, Mat.store(1n, d, g, xs, off))
1864
1865def LayerNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer<LN<d>, TA.Mat<n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1866 Layer{p => x => ln_fwd(n, d, eps, p, x), p => s => dy => g => ln_bwd(n, d, eps, p, s, dy, g), op => p => g => st => ln_step(d, op, p, g, st), u => LN{TA.Mat.zeros(1n, d), TA.Mat.zeros(1n, d)}, p => ln_dump(d, p), p => xs => off => ln_load(d, p, xs, off), p => xs => off => ln_store(d, p, xs, off), u => U32.from_nat(Nat.mul(2n, d))}
1867
1868# =====================================================================================================
1869# Embedding: row i of the output is row ids[i] of the table (v x d); the gradient adds each row back
1870# =====================================================================================================
1871
1872# row i of the output (d numbers) copied from row ids[i] of the table, for the first n ids
1873def emb_cp(left: Nat, ids: List<&2, Nat>, +i: U32, +d: Nat, st: A2) -> A2:
1874 match left ids st:
1875 case 1n+q Con{+id, t} A2{e, o}:
1876 emb_cp(q, t, (i + 1 : U32), d, copy_n(Nat.sub(d, 1n), (U32.from_nat(id) * U32.from_nat(d) : U32), (i * U32.from_nat(d) : U32), o, Array.get(F32, e, (U32.from_nat(id) * U32.from_nat(d) : U32))))
1877 case _ _ st:
1878 st
1879
1880def emb_f2(-n: Nat, -v: Nat, -d: Nat, +ids: List<&2, Nat>, r: A2) -> Fwd<TA.Mat<v, d>, TA.Mat<n, d>, List<&2, Nat>>:
1881 match r:
1882 case A2{e, o}:
1883 Fwd{TA.Mat{e}, TA.Mat{o}, ids}
1884
1885def emb_f(+n: Nat, +v: Nat, +d: Nat, +ids: List<&2, Nat>, e: TA.Mat<v, d>) -> Fwd<TA.Mat<v, d>, TA.Mat<n, d>, List<&2, Nat>>:
1886 match e:
1887 case TA.Mat{ea}:
1888 emb_f2(n, v, d, ids, emb_cp(n, ids, 0, d, A2{ea, zeros_a(Nat.mul(n, d))}))
1889
1890# g[gi + j] = g[gi + j] + dy[di + j], j = 0..k
1891def emb_add(k: Nat, +gi: U32, +di: U32, rg: Array<F32> & F32, rd: Array<F32> & F32) -> A2:
1892 match k rg rd:
1893 case 0n Tuple{g, +a} Tuple{dy, +b}:
1894 A2{Array.set(F32, g, gi, (a + b : F32)), dy}
1895 case 1n+q Tuple{g, +a} Tuple{dy, +b}:
1896 emb_add(q, (gi + 1 : U32), (di + 1 : U32), Array.get(F32, Array.set(F32, g, gi, (a + b : F32)), (gi + 1 : U32)), Array.get(F32, dy, (di + 1 : U32)))
1897
1898# each dY row added into the gradient's row ids[i], in order, by index
1899def emb_sc(left: Nat, ids: List<&2, Nat>, +i: U32, +d: Nat, st: A2) -> A2:
1900 match left ids st:
1901 case 1n+q Con{+id, t} A2{g, dy}:
1902 emb_sc(q, t, (i + 1 : U32), d, emb_add(Nat.sub(d, 1n), (U32.from_nat(id) * U32.from_nat(d) : U32), (i * U32.from_nat(d) : U32), Array.get(F32, g, (U32.from_nat(id) * U32.from_nat(d) : U32)), Array.get(F32, dy, (i * U32.from_nat(d) : U32))))
1903 case _ _ st:
1904 st
1905
1906def emb_b2(-v: Nat, -d: Nat, e: TA.Mat<v, d>, r: A2) -> Bwd<TA.Mat<v, d>, List<&2, Nat>>:
1907 match r:
1908 case A2{g, dy}:
1909 Bwd{e, Nil{}, TA.Mat{g}}
1910
1911def emb_bwd(+n: Nat, +v: Nat, +d: Nat, e: TA.Mat<v, d>, ids: List<&2, Nat>, dy: TA.Mat<n, d>, g: TA.Mat<v, d>) -> Bwd<TA.Mat<v, d>, List<&2, Nat>>:
1912 match dy g:
1913 case TA.Mat{da} TA.Mat{ga}:
1914 emb_b2(v, d, e, emb_sc(n, ids, 0, d, A2{ga, da}))
1915
1916# a table of v rows of d numbers, looked up by n ids; an id has no gradient, so dX is the empty list
1917def Embedding(+n: Nat, +v: Nat, +d: Nat) -> Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
1918 Layer{p => x => emb_f(n, v, d, x, p), p => s => dy => g => emb_bwd(n, v, d, p, s, dy, g), op => p => g => st => mat_step(v, d, op, p, g, st), u => TA.Mat.zeros(v, d), p => Mat.dump(v, d, p), p => xs => off => Mat.restore(v, d, p, xs, off), p => xs => off => Mat.store(v, d, p, xs, off), u => U32.from_nat(Nat.mul(v, d))}
1919
1920# =====================================================================================================
1921# Linear without bias: y = x·Wᵀ (Llama's projections)
1922# =====================================================================================================
1923
1924def lnb_f2(-n: Nat, -i: Nat, -o: Nat, r: TA.MMulNT<n, i, o>) -> Fwd<TA.Mat<o, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
1925 match r:
1926 case TA.MMulNT{x, w, y}:
1927 Fwd{w, y, x}
1928
1929def lnb_b3(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat<o, i>, dx: TA.Mat<n, i>, r: TA.MMulTN<o, n, i>) -> Bwd<TA.Mat<o, i>, TA.Mat<n, i>>:
1930 match r:
1931 case TA.MMulTN{dy, x, g}:
1932 Bwd{w, dx, g}
1933
1934def lnb_b2(+n: Nat, +i: Nat, +o: Nat, +par: Nat, x: TA.Mat<n, i>, g: TA.Mat<o, i>, r: TA.MMul<n, o, i>) -> Bwd<TA.Mat<o, i>, TA.Mat<n, i>>:
1935 match r:
1936 case TA.MMul{dy, w, dx}:
1937 lnb_b3(n, i, o, w, dx, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, g))
1938
1939def LinearNB(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer<TA.Mat<o, i>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
1940 Layer{w => x => lnb_f2(n, i, o, TA.Mat.matmul_nt(n, i, o, par, x, w)), w => s => dy => g => lnb_b2(n, i, o, par, s, g, TA.Mat.matmul_precise(n, o, i, par, dy, w)), op => p => g => st => mat_step(o, i, op, p, g, st), u => TA.Mat.zeros(o, i), p => Mat.dump(o, i, p), p => xs => off => Mat.restore(o, i, p, xs, off), p => xs => off => Mat.store(o, i, p, xs, off), u => U32.from_nat(Nat.mul(o, i))}
1941
1942# =====================================================================================================
1943# Combinators: Id, Fork (one input to two layers), Both (two inputs, two layers), Add2, Mul2, Residual
1944# =====================================================================================================
1945
1946def Id(~X: Type) -> Layer<Unit, X, X, Unit>:
1947 Layer{p => x => Fwd{p, x, Unit{}}, p => s => dy => g => Bwd{p, dy, g}, op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
1948
1949# one optimizer step on two layers' parameters, of any types
1950def pair_step2(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, +op: Opt, p: Two<P1, P2>, g: Two<P1, P2>, st: Two<OptState<P1>, OptState<P2>>) -> Step<Two<P1, P2>>:
1951 match p g st:
1952 case Two{p1, p2} Two{g1, g2} Two{s1, s2}:
1953 pair_s3(P1, P2, Layer.step(P1, X1, Y1, S1, la, op, p1, g1, s1), Layer.step(P2, X2, Y2, S2, lb, op, p2, g2, s2))
1954
1955def pair_step(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, +op: Opt, p: Two<P1, P2>, g: Two<P1, P2>, st: OptState<Two<P1, P2>>) -> Step<Two<P1, P2>>:
1956 pair_step2(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, op, p, g, OptState.split(P1, P2, st))
1957
1958def pair_f4(-P1: Type, -P2: Type, -Y1: Type, -Y2: Type, -S1: Type, -S2: Type, p1: P1, y1: Y1, s1: S1, r: Fwd<P2, Y2, S2>) -> Fwd<Two<P1, P2>, Two<Y1, Y2>, Two<S1, S2>>:
1959 match r:
1960 case Fwd{p2, y2, s2}:
1961 Fwd{Two{p1, p2}, Two{y1, y2}, Two{s1, s2}}
1962
1963# runs lb on x2 once la's forward pass is done
1964def pair_f3(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer<P2, X2, Y2, S2>, -P1: Type, -Y1: Type, -S1: Type, p2: P2, x2: X2, r: Fwd<P1, Y1, S1>) -> Fwd<Two<P1, P2>, Two<Y1, Y2>, Two<S1, S2>>:
1965 match r:
1966 case Fwd{p1, y1, s1}:
1967 pair_f4(P1, P2, Y1, Y2, S1, S2, p1, y1, s1, Layer.fwd(P2, X2, Y2, S2, lb, p2, x2))
1968
1969# ---- Fork: x goes to both layers (copied by `dup`); the two input gradients are summed by `add` ----
1970
1971def fork_f2(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X, Y1, S1>, ~lb: Layer<P2, X, Y2, S2>, p1: P1, p2: P2, d: Two<X, X>) -> Fwd<Two<P1, P2>, Two<Y1, Y2>, Two<S1, S2>>:
1972 match d:
1973 case Two{xa, xb}:
1974 pair_f3(~P2, ~X, ~Y2, ~S2, ~lb, P1, Y1, S1, p2, xb, Layer.fwd(P1, X, Y1, S1, la, p1, xa))
1975
1976def fork_b3(~X: Type, ~add: X -> X -> X, -P1: Type, -P2: Type, p2: P2, g2: P2, dx2: X, r: Bwd<P1, X>) -> Bwd<Two<P1, P2>, X>:
1977 match r:
1978 case Bwd{p1, dx1, g1}:
1979 Bwd{Two{p1, p2}, add(dx1, dx2), Two{g1, g2}}
1980
1981def fork_b2(~P1: Type, ~X: Type, ~Y1: Type, ~S1: Type, ~add: X -> X -> X, ~la: Layer<P1, X, Y1, S1>, -P2: Type, p1: P1, s1: S1, d1: Y1, g1: P1, r: Bwd<P2, X>) -> Bwd<Two<P1, P2>, X>:
1982 match r:
1983 case Bwd{p2, dx2, g2}:
1984 fork_b3(~X, ~add, P1, P2, p2, g2, dx2, Layer.bwd(P1, X, Y1, S1, la, p1, s1, d1, g1))
1985
1986def fork_bwd(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~add: X -> X -> X, ~la: Layer<P1, X, Y1, S1>, ~lb: Layer<P2, X, Y2, S2>, p: Two<P1, P2>, s: Two<S1, S2>, dy: Two<Y1, Y2>, g: Two<P1, P2>) -> Bwd<Two<P1, P2>, X>:
1987 match p s dy g:
1988 case Two{p1, p2} Two{s1, s2} Two{d1, d2} Two{g1, g2}:
1989 fork_b2(~P1, ~X, ~Y1, ~S1, ~add, ~la, P2, p1, s1, d1, g1, Layer.bwd(P2, X, Y2, S2, lb, p2, s2, d2, g2))
1990
1991def fork_fwd(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~dup: X -> Two<X, X>, ~la: Layer<P1, X, Y1, S1>, ~lb: Layer<P2, X, Y2, S2>, p: Two<P1, P2>, x: X) -> Fwd<Two<P1, P2>, Two<Y1, Y2>, Two<S1, S2>>:
1992 match p:
1993 case Two{p1, p2}:
1994 fork_f2(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p1, p2, dup(x))
1995
1996def Fork(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~dup: X -> Two<X, X>, ~add: X -> X -> X, ~la: Layer<P1, X, Y1, S1>, ~lb: Layer<P2, X, Y2, S2>) -> Layer<Two<P1, P2>, X, Two<Y1, Y2>, Two<S1, S2>>:
1997 Layer{p => x => fork_fwd(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~dup, ~la, ~lb, p, x), p => s => dy => g => fork_bwd(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~add, ~la, ~lb, p, s, dy, g), op => p => g => st => pair_step(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P1, X, Y1, S1, la), Layer.zeros(P2, X, Y2, S2, lb)}, p => pair_dump(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X, Y1, S1, la) + Layer.count(P2, X, Y2, S2, lb) : U32)}
1998
1999# ---- Both: la on the first of a pair, lb on the second ----
2000
2001def both_b3(-P1: Type, -P2: Type, -X1: Type, -X2: Type, p2: P2, g2: P2, dx2: X2, r: Bwd<P1, X1>) -> Bwd<Two<P1, P2>, Two<X1, X2>>:
2002 match r:
2003 case Bwd{p1, dx1, g1}:
2004 Bwd{Two{p1, p2}, Two{dx1, dx2}, Two{g1, g2}}
2005
2006def both_b2(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~la: Layer<P1, X1, Y1, S1>, -P2: Type, -X2: Type, p1: P1, s1: S1, d1: Y1, g1: P1, r: Bwd<P2, X2>) -> Bwd<Two<P1, P2>, Two<X1, X2>>:
2007 match r:
2008 case Bwd{p2, dx2, g2}:
2009 both_b3(P1, P2, X1, X2, p2, g2, dx2, Layer.bwd(P1, X1, Y1, S1, la, p1, s1, d1, g1))
2010
2011def both_bwd(~P1: Type, ~P2: Type, ~X1: Type, ~X2: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, p: Two<P1, P2>, s: Two<S1, S2>, dy: Two<Y1, Y2>, g: Two<P1, P2>) -> Bwd<Two<P1, P2>, Two<X1, X2>>:
2012 match p s dy g:
2013 case Two{p1, p2} Two{s1, s2} Two{d1, d2} Two{g1, g2}:
2014 both_b2(~P1, ~X1, ~Y1, ~S1, ~la, P2, X2, p1, s1, d1, g1, Layer.bwd(P2, X2, Y2, S2, lb, p2, s2, d2, g2))
2015
2016def both_fwd(~P1: Type, ~P2: Type, ~X1: Type, ~X2: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>, p: Two<P1, P2>, x: Two<X1, X2>) -> Fwd<Two<P1, P2>, Two<Y1, Y2>, Two<S1, S2>>:
2017 match p x:
2018 case Two{p1, p2} Two{x1, x2}:
2019 pair_f3(~P2, ~X2, ~Y2, ~S2, ~lb, P1, Y1, S1, p2, x2, Layer.fwd(P1, X1, Y1, S1, la, p1, x1))
2020
2021def Both(~P1: Type, ~P2: Type, ~X1: Type, ~X2: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer<P1, X1, Y1, S1>, ~lb: Layer<P2, X2, Y2, S2>) -> Layer<Two<P1, P2>, Two<X1, X2>, Two<Y1, Y2>, Two<S1, S2>>:
2022 Layer{p => x => both_fwd(~P1, ~P2, ~X1, ~X2, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p, x), p => s => dy => g => both_bwd(~P1, ~P2, ~X1, ~X2, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p, s, dy, g), op => p => g => st => pair_step(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P1, X1, Y1, S1, la), Layer.zeros(P2, X2, Y2, S2, lb)}, p => pair_dump(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X1, Y1, S1, la) + Layer.count(P2, X2, Y2, S2, lb) : U32)}
2023
2024# ---- Tie: first, mid, last, where first and last share one parameter (tied embeddings: the table is also the output
2025# layer's weight). Forward runs the three in order with the same P for first and last; backward runs last, mid, first,
2026# both adding into the same gradient, as PyTorch does for a tensor used twice. ----
2027
2028type Three<-A: Type, -B: Type, -C: Type> is Type:
2029 Three{a: A, b: B, c: C}
2030
2031def tie_f4(-P: Type, -Q: Type, -W: Type, -S1: Type, -S2: Type, -S3: Type, q: Q, s1: S1, s2: S2, r: Fwd<P, W, S3>) -> Fwd<Two<P, Q>, W, Three<S1, S2, S3>>:
2032 match r:
2033 case Fwd{p, w, s3}:
2034 Fwd{Two{p, q}, w, Three{s1, s2, s3}}
2035
2036def tie_f3(~P: Type, ~Z: Type, ~W: Type, ~S3: Type, ~lc: Layer<P, Z, W, S3>, -Q: Type, -S1: Type, -S2: Type, p: P, s1: S1, r: Fwd<Q, Z, S2>) -> Fwd<Two<P, Q>, W, Three<S1, S2, S3>>:
2037 match r:
2038 case Fwd{q, z, s2}:
2039 tie_f4(P, Q, W, S1, S2, S3, q, s1, s2, Layer.fwd(P, Z, W, S3, lc, p, z))
2040
2041def tie_f2(~P: Type, ~Q: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S2: Type, ~S3: Type, ~lb: Layer<Q, Y, Z, S2>, ~lc: Layer<P, Z, W, S3>, -S1: Type, q: Q, r: Fwd<P, Y, S1>) -> Fwd<Two<P, Q>, W, Three<S1, S2, S3>>:
2042 match r:
2043 case Fwd{p, y, s1}:
2044 tie_f3(~P, ~Z, ~W, ~S3, ~lc, Q, S1, S2, p, s1, Layer.fwd(Q, Y, Z, S2, lb, q, y))
2045
2046def tie_fwd(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S1: Type, ~S2: Type, ~S3: Type, ~la: Layer<P, X, Y, S1>, ~lb: Layer<Q, Y, Z, S2>, ~lc: Layer<P, Z, W, S3>, pq: Two<P, Q>, x: X) -> Fwd<Two<P, Q>, W, Three<S1, S2, S3>>:
2047 match pq:
2048 case Two{p, q}:
2049 tie_f2(~P, ~Q, ~Y, ~Z, ~W, ~S2, ~S3, ~lb, ~lc, S1, q, Layer.fwd(P, X, Y, S1, la, p, x))
2050
2051def tie_b4(-P: Type, -Q: Type, -X: Type, q: Q, gq: Q, r: Bwd<P, X>) -> Bwd<Two<P, Q>, X>:
2052 match r:
2053 case Bwd{p, dx, gp}:
2054 Bwd{Two{p, q}, dx, Two{gp, gq}}
2055
2056def tie_b3(~P: Type, ~X: Type, ~Y: Type, ~S1: Type, ~la: Layer<P, X, Y, S1>, -Q: Type, p: P, s1: S1, gp: P, r: Bwd<Q, Y>) -> Bwd<Two<P, Q>, X>:
2057 match r:
2058 case Bwd{q, dy, gq}:
2059 tie_b4(P, Q, X, q, gq, Layer.bwd(P, X, Y, S1, la, p, s1, dy, gp))
2060
2061def tie_b2(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer<P, X, Y, S1>, ~lb: Layer<Q, Y, Z, S2>, q: Q, s1: S1, s2: S2, gq: Q, r: Bwd<P, Z>) -> Bwd<Two<P, Q>, X>:
2062 match r:
2063 case Bwd{p, dz, gp}:
2064 tie_b3(~P, ~X, ~Y, ~S1, ~la, Q, p, s1, gp, Layer.bwd(Q, Y, Z, S2, lb, q, s2, dz, gq))
2065
2066def tie_bwd(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S1: Type, ~S2: Type, ~S3: Type, ~la: Layer<P, X, Y, S1>, ~lb: Layer<Q, Y, Z, S2>, ~lc: Layer<P, Z, W, S3>, pq: Two<P, Q>, s: Three<S1, S2, S3>, dw: W, g: Two<P, Q>) -> Bwd<Two<P, Q>, X>:
2067 match pq s g:
2068 case Two{p, q} Three{s1, s2, s3} Two{gp, gq}:
2069 tie_b2(~P, ~Q, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, q, s1, s2, gq, Layer.bwd(P, Z, W, S3, lc, p, s3, dw, gp))
2070
2071# The shared parameter steps once (through first's step; last's must be the same kind of layer for P, e.g. Embedding
2072# and LinearNB over the same v x d matrix), the middle with its own.
2073def Tie(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S1: Type, ~S2: Type, ~S3: Type, ~la: Layer<P, X, Y, S1>, ~lb: Layer<Q, Y, Z, S2>, ~lc: Layer<P, Z, W, S3>) -> Layer<Two<P, Q>, X, W, Three<S1, S2, S3>>:
2074 Layer{pq => x => tie_fwd(~P, ~Q, ~X, ~Y, ~Z, ~W, ~S1, ~S2, ~S3, ~la, ~lb, ~lc, pq, x), pq => s => dw => g => tie_bwd(~P, ~Q, ~X, ~Y, ~Z, ~W, ~S1, ~S2, ~S3, ~la, ~lb, ~lc, pq, s, dw, g), op => p => g => st => pair_step(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P, X, Y, S1, la), Layer.zeros(Q, Y, Z, S2, lb)}, p => pair_dump(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P, X, Y, S1, la) + Layer.count(Q, Y, Z, S2, lb) : U32)}
2075
2076# ---- Add2 and Mul2: element-wise sum and product of a pair of matrices ----
2077
2078def add2_f(+n: Nat, +c: Nat, x: Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> Fwd<Unit, TA.Mat<n, c>, Unit>:
2079 match x:
2080 case Two{a, b}:
2081 Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}}
2082
2083def Add2(+n: Nat, +c: Nat) -> Layer<Unit, Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Unit>:
2084 Layer{p => x => add2_f(n, c, x), p => s => dy => g => Bwd{p, Mat.dup(n, c, dy), g}, op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
2085
2086def mul2_f3(-n: Nat, -c: Nat, a2: TA.Mat<n, c>, b2: TA.Mat<n, c>, r: TA.MPair<n, c>) -> Fwd<Unit, TA.Mat<n, c>, Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
2087 match r:
2088 case TA.MPair{y, b}:
2089 Fwd{Unit{}, y, Two{a2, b2}}
2090
2091def mul2_f2(+n: Nat, +c: Nat, da: Two<TA.Mat<n, c>, TA.Mat<n, c>>, db: Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> Fwd<Unit, TA.Mat<n, c>, Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
2092 match da db:
2093 case Two{a1, a2} Two{b1, b2}:
2094 mul2_f3(n, c, a2, b2, Ew.zip(~f32_mul, n, c, a1, b1))
2095
2096def mul2_f(+n: Nat, +c: Nat, x: Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> Fwd<Unit, TA.Mat<n, c>, Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
2097 match x:
2098 case Two{a, b}:
2099 mul2_f2(n, c, Mat.dup(n, c, a), Mat.dup(n, c, b))
2100
2101def mpair_a(-n: Nat, -c: Nat, r: TA.MPair<n, c>) -> TA.Mat<n, c>:
2102 match r:
2103 case TA.MPair{a, b}:
2104 a
2105
2106# d(a*b) = (dy*b, dy*a)
2107def mul2_b2(+n: Nat, +c: Nat, s: Two<TA.Mat<n, c>, TA.Mat<n, c>>, d: Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> Bwd<Unit, Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
2108 match s d:
2109 case Two{a, b} Two{d1, d2}:
2110 Bwd{Unit{}, Two{mpair_a(n, c, Ew.zip(~f32_mul, n, c, d1, b)), mpair_a(n, c, Ew.zip(~f32_mul, n, c, d2, a))}, Unit{}}
2111
2112def Mul2(+n: Nat, +c: Nat) -> Layer<Unit, Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
2113 Layer{p => x => mul2_f(n, c, x), p => s => dy => g => mul2_b2(n, c, s, Mat.dup(n, c, dy)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
2114
2115# ---- Residual: y = x + f(x), for any layer f from n x c to n x c ----
2116
2117def Residual(~n: Nat, ~c: Nat, ~P: Type, ~S: Type, ~f: Layer<P, TA.Mat<n, c>, TA.Mat<n, c>, S>) -> Layer<Two<Two<Unit, P>, Unit>, TA.Mat<n, c>, TA.Mat<n, c>, Two<Two<Unit, S>, Unit>>:
2118 Seq(~Two<Unit, P>, ~Unit, ~TA.Mat<n, c>, ~Two<TA.Mat<n, c>, TA.Mat<n, c>>, ~TA.Mat<n, c>, ~Two<Unit, S>, ~Unit, ~Fork(~Unit, ~P, ~TA.Mat<n, c>, ~TA.Mat<n, c>, ~TA.Mat<n, c>, ~Unit, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~Id(~TA.Mat<n, c>), ~f), ~Add2(n, c))
2119
2120# =====================================================================================================
2121# RoPE (the rotate-half form of Llama and `transformers`): rows are positions (row r is position r mod t), each row
2122# holds h heads of 2*half numbers, and the pair (x[c], x[c + half]) of a head is rotated by pos * inv[c], with
2123# inv[c] = 1 / base^(2c / (2*half)). No parameters, nothing saved: the backward pass rotates by the opposite angle.
2124# =====================================================================================================
2125
2126# Over the arrays: a table of cos and sin for each position mod t and each c (t x half), computed once a call with the
2127# operations the list version used, then every pair rotated in place, so the bits are the list version's.
2128
2129type CS is Type:
2130 CS{c: Array<F32>, s: Array<F32>}
2131
2132# the angle of table entry e: inv[c] * position, position e / half, c = e mod half, inv[c] = 1 / base^(2c / (2*half))
2133def rope_ang(+e: U32, +half: U32, +base: F32) -> F32:
2134 ((1.0 / F32.pow(base, (U32.to_f32((U32.mod(e, half) + U32.mod(e, half) : U32)) / U32.to_f32((half + half : U32)) : F32)) : F32) * U32.to_f32(U32.div(e, half)) : F32)
2135
2136# entries e.. of the tables, left more after e
2137def rope_tab(left: Nat, +e: U32, +half: U32, +base: F32, tb: CS) -> CS:
2138 match left tb:
2139 case 0n CS{cs, sn}:
2140 +ang = rope_ang(e, half, base)
2141 CS{Array.set(F32, cs, e, F32.cos(ang)), Array.set(F32, sn, e, F32.sin(ang))}
2142 case 1n+q CS{cs, sn}:
2143 +ang = rope_ang(e, half, base)
2144 rope_tab(q, (e + 1 : U32), half, base, CS{Array.set(F32, cs, e, F32.cos(ang)), Array.set(F32, sn, e, F32.sin(ang))})
2145
2146type RS3 is Type:
2147 RS3{x: Array<F32>, c: Array<F32>, s: Array<F32>}
2148
2149# one pair (a = x[i1], b = x[i2]) rotated in place; fwd picks the direction: (a cos - b sin, b cos + a sin) forward,
2150# (a cos + b sin, b cos - a sin) backward, in PyTorch's order of operations for each
2151def rope_p3(fwd: Bool, +i1: U32, +i2: U32, +a: F32, r2: Array<F32> & F32, rc: Array<F32> & F32, rs: Array<F32> & F32) -> RS3:
2152 match fwd r2 rc rs:
2153 case True{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}:
2154 RS3{Array.set(F32, Array.set(F32, x, i1, ((a * c : F32) + (F32.neg(b) * z : F32) : F32)), i2, ((b * c : F32) + (a * z : F32) : F32)), cs, sn}
2155 case False{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}:
2156 RS3{Array.set(F32, Array.set(F32, x, i1, ((a * c : F32) + (b * z : F32) : F32)), i2, ((b * c : F32) + F32.neg((a * z : F32)) : F32)), cs, sn}
2157
2158def rope_p2(+fwd: Bool, +i1: U32, +i2: U32, +k: U32, cs: Array<F32>, sn: Array<F32>, r1: Array<F32> & F32) -> RS3:
2159 match r1:
2160 case Tuple{x, +a}:
2161 rope_p3(fwd, i1, i2, a, Array.get(F32, x, i2), Array.get(F32, cs, k), Array.get(F32, sn, k))
2162
2163# pair e of the n x w matrix: row e / hh (hh = h * half pairs a row), head (e mod hh) / half, column c of it
2164def rope_one(+e: U32, +fwd: Bool, +hh: U32, +half: U32, +w: U32, +t: U32, st: RS3) -> RS3:
2165 match st:
2166 case RS3{x, cs, sn}:
2167 +r = U32.div(e, hh)
2168 +c = U32.mod(U32.mod(e, hh), half)
2169 +i1 = (((r * w : U32) + (U32.div(U32.mod(e, hh), half) * (half + half : U32) : U32) : U32) + c : U32)
2170 rope_p2(fwd, i1, (i1 + half : U32), ((U32.mod(r, t) * half : U32) + c : U32), cs, sn, Array.get(F32, x, i1))
2171
2172# pairs e.., left more after e
2173def rope_e(left: Nat, +e: U32, +fwd: Bool, +hh: U32, +half: U32, +w: U32, +t: U32, st: RS3) -> RS3:
2174 match left:
2175 case 0n:
2176 rope_one(e, fwd, hh, half, w, t, st)
2177 case 1n+q:
2178 rope_e(q, (e + 1 : U32), fwd, hh, half, w, t, rope_one(e, fwd, hh, half, w, t, st))
2179
2180def rope_fin(-n: Nat, -w: Nat, st: RS3) -> TA.Mat<n, w>:
2181 match st:
2182 case RS3{x, cs, sn}:
2183 TA.Mat{x}
2184
2185def rope_go2(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, x: Array<F32>, tb: CS) -> TA.Mat<n, w>:
2186 match tb:
2187 case CS{cs, sn}:
2188 rope_fin(n, w, rope_e(Nat.sub(Nat.mul(n, Nat.mul(h, half)), 1n), 0, fwd, U32.from_nat(Nat.mul(h, half)), U32.from_nat(half), U32.from_nat(w), U32.from_nat(t), RS3{x, cs, sn}))
2189
2190def rope_go(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, +base: F32, m: TA.Mat<n, w>) -> TA.Mat<n, w>:
2191 match m:
2192 case TA.Mat{x}:
2193 rope_go2(n, w, fwd, t, h, half, x, rope_tab(Nat.sub(Nat.mul(t, half), 1n), 0, U32.from_nat(half), base, CS{zeros_a(Nat.mul(t, half)), zeros_a(Nat.mul(t, half))}))
2194
2195# n rows (b sequences of t positions), w = h heads of 2*half numbers each
2196def RoPE(+n: Nat, +w: Nat, +t: Nat, +h: Nat, +half: Nat, +base: F32, -e: {w == Nat.mul(h, Nat.add(half, half)) : Nat}) -> Layer<Unit, TA.Mat<n, w>, TA.Mat<n, w>, Unit>:
2197 Layer{p => x => Fwd{p, rope_go(n, w, True{}, t, h, half, base, x), Unit{}}, p => s => dy => g => Bwd{p, rope_go(n, w, False{}, t, h, half, base, dy), g}, op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
2198
2199# =====================================================================================================
2200# Causal attention with grouped-query heads: queries n x (h*hd), keys and values n x (kv*hd), h = kv*g; query head
2201# i reads key/value head i / g. Rows are b sequences of t positions; position p sees positions 0..p of its sequence.
2202# S = (Q·Kᵀ) * scale, P = softmax(S) over the visible positions, O = P·V
2203# The forward pass keeps Q, K, V; the backward pass computes P again from them (recomputation, not storage).
2204# =====================================================================================================
2205
2206# Q, K and V stay in their arrays. Each sequence's t rows are copied out once into arrays of their own, and every
2207# head's probabilities P (t x t) are an array too. Each sum is added in the order the earlier list version added it
2208# (left to right from +0.0), and each row of P is softmaxed in place over the row's visible positions (prow), so the
2209# bits are the list version's.
2210
2211type AttnS is Type:
2212 AttnS{q: Array<F32>, k: Array<F32>, v: Array<F32>}
2213
2214# the key/value head that query head hi reads, with g query heads per key/value head; the forward and the backward pass
2215# both pick it here (the law kv_of_group says which heads share one)
2216def kv_of(+hi: Nat, +g: Nat) -> Nat:
2217 Nat.div(hi, g)
2218
2219# xs reversed onto acc
2220def rcat(xs: List<&2, F32>, acc: List<&2, F32>) -> List<&2, F32>:
2221 match xs:
2222 case Nil{}:
2223 acc
2224 case Con{+x, t}:
2225 rcat(t, x <> acc)
2226
2227type DR is Type:
2228 DR{a: Array<F32>, b: Array<F32>, row: List<&2, F32>}
2229
2230# the dots of one row of a (hd numbers from ia) with consecutive rows of b (hd numbers from ib, each bs further): the
2231# current one r and m more, in order
2232def drow(m: Nat, +hd: Nat, +ia: U32, +ib: U32, +bs: U32, acc: List<&2, F32>, r: D2) -> DR:
2233 match m r:
2234 case 0n D2{a, b, +s}:
2235 DR{a, b, rcat(s <> acc, Nil{})}
2236 case 1n+p D2{a, b, +s}:
2237 drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, dot0(hd, ia, (ib + bs : U32), 1, 1, a, b))
2238
2239def dr_row(r: DR) -> List<&2, F32>:
2240 match r:
2241 case DR{a, b, row}:
2242 row
2243
2244# the scores of the query at position i (row i of q, qs apart, the head's hd columns from qc) against keys 0..i (rows
2245# of k, ks apart, from kc): the one place attention reads keys, i + 1 of them (the law causal_keys)
2246def qrow(+i: Nat, +hd: Nat, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>) -> DR:
2247 drow(i, hd, ((U32.from_nat(i) * qs : U32) + qc : U32), kc, ks, Nil{}, dot0(hd, ((U32.from_nat(i) * qs : U32) + qc : U32), kc, 1, 1, q, k))
2248
2249type QK is Type:
2250 QK{q: Array<F32>, k: Array<F32>, p: Array<F32>}
2251
2252type PW is Type:
2253 PW{p: Array<F32>, n: Nat}
2254
2255# xs scaled by sc, written into p from idx on; n counts the numbers written (from the n given)
2256def pw(xs: List<&2, F32>, +idx: U32, +sc: F32, n: Nat, p: Array<F32>) -> PW:
2257 match xs:
2258 case Nil{}:
2259 PW{p, n}
2260 case Con{x, t}:
2261 pw(t, (idx + 1 : U32), sc, 1n+n, Array.set(F32, p, idx, (x * sc : F32)))
2262
2263type PR is Type:
2264 PR{q: Array<F32>, k: Array<F32>, p: Array<F32>, n: Nat}
2265
2266def pr_n(r: PR) -> Nat:
2267 match r:
2268 case PR{q, k, p, n}:
2269 n
2270
2271def pr_of(q: Array<F32>, k: Array<F32>, r: PC) -> PR:
2272 match r:
2273 case PC{p, c}:
2274 PR{q, k, p, c}
2275
2276# the softmax over the n numbers pw wrote; the row's count is the one the softmax reports
2277def prow3(+off: U32, q: Array<F32>, k: Array<F32>, w: PW) -> PR:
2278 match w:
2279 case PW{p, +n}:
2280 pr_of(q, k, psmc(n, off, p))
2281
2282def prow2(+sc: F32, +off: U32, p: Array<F32>, r: DR) -> PR:
2283 match r:
2284 case DR{q, k, row}:
2285 prow3(off, q, k, pw(row, off, sc, 0n, p))
2286
2287# row i of a head's P (t x t, row-major, created zeros): the scaled scores of query i against keys 0..i (qrow) written
2288# from the row's start, then the softmax over as many positions as were written: i + 1 (the law causal). The positions
2289# after i are never written and keep their zeros.
2290def prow(+i: Nat, +t: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>, p: Array<F32>) -> PR:
2291 prow2(sc, (U32.from_nat(i) * U32.from_nat(t) : U32), p, qrow(i, hd, qs, qc, ks, kc, q, k))
2292
2293# rows i.. of a head's P: r holds rows up to i; left rows more after it. Queries are rows of q (qs apart, the head's
2294# columns from qc), keys rows of k (ks apart, from kc).
2295def prows(left: Nat, +i: Nat, +t: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, r: PR) -> QK:
2296 match left r:
2297 case 0n PR{q, k, p, n}:
2298 QK{q, k, p}
2299 case 1n+m PR{q, k, p, n}:
2300 prows(m, 1n+i, t, hd, sc, qs, qc, ks, kc, prow(1n+i, t, hd, sc, qs, qc, ks, kc, q, k, p))
2301
2302# P of one head: queries in q (row stride qs, columns from qc), keys in k (row stride ks, columns from kc)
2303def probs(+t: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>) -> QK:
2304 prows(Nat.sub(t, 1n), 0n, t, hd, sc, qs, qc, ks, kc, prow(0n, t, hd, sc, qs, qc, ks, kc, q, k, zeros_a(Nat.mul(t, t))))
2305
2306type PO is Type:
2307 PO{a: Array<F32>, b: Array<F32>, o: Array<F32>}
2308
2309# o[i·ow + oc + d] = a[i·t + 0]·b[0·bs + bc + d] + ... + a[i·t + t-1]·b[(t-1)·bs + bc + d], for idx = i·c + d from idx
2310# on, left more after r (P·V, and dQ = G·K): each sum left to right from +0.0
2311def pmul(left: Nat, +idx: U32, +c: U32, +t: Nat, +bs: U32, +bc: U32, +ow: U32, +oc: U32, o: Array<F32>, r: D2) -> PO:
2312 match left r:
2313 case 0n D2{a, b, +s}:
2314 PO{a, b, Array.set(F32, o, ((U32.div(idx, c) * ow : U32) + (oc + U32.mod(idx, c) : U32) : U32), s)}
2315 case 1n+n D2{a, b, +s}:
2316 pmul(n, (idx + 1 : U32), c, t, bs, bc, ow, oc, Array.set(F32, o, ((U32.div(idx, c) * ow : U32) + (oc + U32.mod(idx, c) : U32) : U32), s), dot0(t, (U32.div((idx + 1 : U32), c) * U32.from_nat(t) : U32), (bc + U32.mod((idx + 1 : U32), c) : U32), 1, bs, a, b))
2317
2318def acc_set(+idx: U32, +s: F32, r: Array<F32> & F32) -> Array<F32>:
2319 match r:
2320 case Tuple{o, +old}:
2321 Array.set(F32, o, idx, (old + s : F32))
2322
2323# o[j·c + d] += a[0·t + j]·b[0·bs + bc + d] + ... + a[(t-1)·t + j]·b[(t-1)·bs + bc + d], for idx = j·c + d from idx on
2324# (dK += Gᵀ·Q and dV += Pᵀ·dO): each sum left to right from +0.0, then added into the group's running sum
2325def pacc(left: Nat, +idx: U32, +c: U32, +t: Nat, +bs: U32, +bc: U32, o: Array<F32>, r: D2) -> PO:
2326 match left r:
2327 case 0n D2{a, b, +s}:
2328 PO{a, b, acc_set(idx, s, Array.get(F32, o, idx))}
2329 case 1n+n D2{a, b, +s}:
2330 pacc(n, (idx + 1 : U32), c, t, bs, bc, acc_set(idx, s, Array.get(F32, o, idx)), dot0(t, U32.div((idx + 1 : U32), c), (bc + U32.mod((idx + 1 : U32), c) : U32), U32.from_nat(t), bs, a, b))
2331
2332# rows r0.. (t of them, c wide) of src into a new array, src handed back
2333def rows_out(+t: Nat, +c: Nat, +r0: Nat, src: Array<F32>) -> A2:
2334 copy_n(Nat.sub(Nat.mul(t, c), 1n), U32.from_nat(Nat.mul(r0, c)), 0, zeros_a(Nat.mul(t, c)), Array.get(F32, src, U32.from_nat(Nat.mul(r0, c))))
2335
2336def a2_b(r: A2) -> Array<F32>:
2337 match r:
2338 case A2{a, b}:
2339 b
2340
2341# a block of t rows, c wide, written into dst from row r0; the block is dropped
2342def rows_in(+t: Nat, +c: Nat, +r0: Nat, blk: Array<F32>, dst: Array<F32>) -> Array<F32>:
2343 a2_b(copy_n(Nat.sub(Nat.mul(t, c), 1n), 0, U32.from_nat(Nat.mul(r0, c)), dst, Array.get(F32, blk, 0)))
2344
2345def pmul_copy(left: Nat, +idx: U32, +c: U32, +ow: U32, +oc: U32, o: Array<F32>, r: Array<F32> & F32) -> A2:
2346 match left r:
2347 case 0n Tuple{s, +x}:
2348 A2{s, Array.set(F32, o, ((U32.div(idx, c) * ow : U32) + (oc + U32.mod(idx, c) : U32) : U32), x)}
2349 case 1n+n Tuple{s, +x}:
2350 pmul_copy(n, (idx + 1 : U32), c, ow, oc, Array.set(F32, o, ((U32.div(idx, c) * ow : U32) + (oc + U32.mod(idx, c) : U32) : U32), x), Array.get(F32, s, (idx + 1 : U32)))
2351
2352# dst's columns oc.. (c of them, ow wide) of rows 0..t-1 from src (t x c): src dropped
2353def cols_in(+t: Nat, +c: Nat, +ow: U32, +oc: U32, src: Array<F32>, dst: Array<F32>) -> Array<F32>:
2354 a2_b(pmul_copy(Nat.sub(Nat.mul(t, c), 1n), 0, U32.from_nat(c), ow, oc, dst, Array.get(F32, src, 0)))
2355
2356# ---- the forward pass of one sequence: every head's O = P·V into its columns of o ----
2357
2358type FS is Type:
2359 FS{q: Array<F32>, k: Array<F32>, v: Array<F32>, o: Array<F32>}
2360
2361type Dm is Data:
2362 Dm{t: Nat, g: Nat, hd: Nat, w: Nat, kvw: Nat, sc: F32}
2363
2364# the columns query head hi reads: its own hd columns of q (hc) and those of its key/value head kv_of(hi, g) in k and v
2365# (kc); both passes take their columns from here (the law kv_cols)
2366type HC is Data:
2367 HC{hc: U32, kc: U32}
2368
2369def head_cols(+dm: Dm, +hi: Nat) -> HC:
2370 match dm:
2371 case Dm{+t, +g, +hd, +w, +kvw, +sc}:
2372 HC{U32.from_nat(Nat.mul(hi, hd)), U32.from_nat(Nat.mul(kv_of(hi, g), hd))}
2373
2374def hc_kc(+c: HC) -> U32:
2375 match c:
2376 case HC{+hc, +kc}:
2377 kc
2378
2379def fs_fin(q: Array<F32>, k: Array<F32>, r: PO) -> FS:
2380 match r:
2381 case PO{p, v, o}:
2382 FS{q, k, v, o}
2383
2384def head_f2(+t: Nat, +hd: Nat, +w: Nat, +kvw: Nat, +hc: U32, +kc: U32, v: Array<F32>, o: Array<F32>, r: QK) -> FS:
2385 match r:
2386 case QK{q, k, p}:
2387 fs_fin(q, k, pmul(Nat.sub(Nat.mul(t, hd), 1n), 0, U32.from_nat(hd), t, U32.from_nat(kvw), kc, U32.from_nat(w), hc, o, dot0(t, 0, kc, 1, U32.from_nat(kvw), p, v)))
2388
2389# one head's forward pass at the columns hc (its queries) and kc (its key/value head) that head_cols gives; it has no
2390# head index, so it reads no other columns
2391def head_fwd(+dm: Dm, +cols: HC, s: FS) -> FS:
2392 match dm cols s:
2393 case Dm{+t, +g, +hd, +w, +kvw, +sc} HC{+hc, +kc} FS{q, k, v, o}:
2394 head_f2(t, hd, w, kvw, hc, kc, v, o, probs(t, hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, q, k))
2395
2396def fheads(left: Nat, +hi: Nat, +dm: Dm, s: FS) -> FS:
2397 match left:
2398 case 0n:
2399 s
2400 case 1n+n:
2401 fheads(n, 1n+hi, dm, head_fwd(dm, head_cols(dm, hi), s))
2402
2403def fs_o(s: FS) -> Array<F32>:
2404 match s:
2405 case FS{q, k, v, o}:
2406 o
2407
2408# ---- the backward pass of one sequence: per key/value head, its g query heads' dQ into their columns of dq, their dK
2409# and dV summed in head order (from +0.0), then copied into the head's columns ----
2410
2411type BS is Type:
2412 BS{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, dq: Array<F32>, dka: Array<F32>, dva: Array<F32>}
2413
2414# g[j] = (p[j] * (g[j] - s)) * sc for j = idx..idx + k; rg and rp hold g[idx] and p[idx]
2415def gfix(k: Nat, +idx: U32, +s: F32, +sc: F32, rg: Array<F32> & F32, rp: Array<F32> & F32) -> A2:
2416 match k rg rp:
2417 case 0n Tuple{g, +d} Tuple{p, +y}:
2418 A2{p, Array.set(F32, g, idx, ((y * (d - s : F32) : F32) * sc : F32))}
2419 case 1n+q Tuple{g, +d} Tuple{p, +y}:
2420 gfix(q, (idx + 1 : U32), s, sc, Array.get(F32, Array.set(F32, g, idx, ((y * (d - s : F32) : F32) * sc : F32)), (idx + 1 : U32)), Array.get(F32, p, (idx + 1 : U32)))
2421
2422def g_set2(+t: Nat, +sc: F32, +off: U32, r: D2) -> A2:
2423 match r:
2424 case D2{g, p, +s}:
2425 gfix(Nat.sub(t, 1n), off, s, sc, Array.get(F32, g, off), Array.get(F32, p, off))
2426
2427type GW is Type:
2428 GW{d: Array<F32>, v: Array<F32>, p: Array<F32>, g: Array<F32>}
2429
2430type GN is Type:
2431 GN{g: Array<F32>, n: Nat}
2432
2433# xs written into g from idx, counted onto n
2434def gw(xs: List<&2, F32>, +idx: U32, n: Nat, g: Array<F32>) -> GN:
2435 match xs:
2436 case Nil{}:
2437 GN{g, n}
2438 case Con{x, t}:
2439 gw(t, (idx + 1 : U32), 1n+n, Array.set(F32, g, idx, x))
2440
2441type GR is Type:
2442 GR{d: Array<F32>, v: Array<F32>, pg: A2, n: Nat}
2443
2444def gr_n(r: GR) -> Nat:
2445 match r:
2446 case GR{d, v, pg, n}:
2447 n
2448
2449# the n positions of the row written: s = sum(dP * P) over them (from +0.0, in order), then G = (P * (dP - s)) * scale
2450def gr3(+sc: F32, +off: U32, d: Array<F32>, v: Array<F32>, p: Array<F32>, w: GN) -> GR:
2451 match w:
2452 case GN{g, +n}:
2453 GR{d, v, g_set2(n, sc, off, dot0(n, off, off, 1, 1, g, p)), n}
2454
2455def gr2(+sc: F32, +off: U32, p: Array<F32>, g: Array<F32>, r: DR) -> GR:
2456 match r:
2457 case DR{d, v, row}:
2458 gr3(sc, off, d, v, p, gw(row, off, 0n, g))
2459
2460# row i of G (from i·t, t wide; softmax's backward pass, then the scale): dP's row against values 0..i only (drow, as
2461# qrow reads keys 0..i), written from the row's start, and s and G over as many positions as were written: i + 1 (the
2462# law causal_grad). G starts as zeros and nothing else writes the row, so its positions after i stay 0: no gradient
2463# reaches a key or value after the query's position. (P is 0 there too, so the full row gave the same G.)
2464def grow(+i: Nat, +t: Nat, +hd: Nat, +sc: F32, +ds: U32, +dc: U32, +vs: U32, +vc: U32, d: Array<F32>, v: Array<F32>, p: Array<F32>, g: Array<F32>) -> GR:
2465 gr2(sc, (U32.from_nat(i) * U32.from_nat(t) : U32), p, g, drow(i, hd, ((U32.from_nat(i) * ds : U32) + dc : U32), vc, vs, Nil{}, dot0(hd, ((U32.from_nat(i) * ds : U32) + dc : U32), vc, 1, 1, d, v)))
2466
2467# rows i.. of G; r is row i
2468def grows(left: Nat, +i: Nat, +t: Nat, +hd: Nat, +sc: F32, +ds: U32, +dc: U32, +vs: U32, +vc: U32, r: GR) -> GW:
2469 match left r:
2470 case 0n GR{d, v, A2{p, g}, n}:
2471 GW{d, v, p, g}
2472 case 1n+m GR{d, v, A2{p, g}, n}:
2473 grows(m, 1n+i, t, hd, sc, ds, dc, vs, vc, grow(1n+i, t, hd, sc, ds, dc, vs, vc, d, v, p, g))
2474
2475def hb5(q: Array<F32>, k: Array<F32>, v: Array<F32>, dq: Array<F32>, dka: Array<F32>, r: PO) -> BS:
2476 match r:
2477 case PO{p, d, dva}:
2478 BS{q, k, v, d, dq, dka, dva}
2479
2480# dK += Gᵀ·Q, then dV += Pᵀ·dO
2481def hb4(+t: Nat, +hd: Nat, +w: Nat, +hc: U32, k: Array<F32>, d: Array<F32>, v: Array<F32>, p: Array<F32>, dq: Array<F32>, dva: Array<F32>, r: PO) -> BS:
2482 match r:
2483 case PO{g, q, dka}:
2484 hb5(q, k, v, dq, dka, pacc(Nat.sub(Nat.mul(t, hd), 1n), 0, U32.from_nat(hd), t, U32.from_nat(w), hc, dva, dot0(t, 0, hc, U32.from_nat(t), U32.from_nat(w), p, d)))
2485
2486# dQ = G·K into the head's columns, then dK
2487def hb3(+t: Nat, +hd: Nat, +w: Nat, +hc: U32, q: Array<F32>, d: Array<F32>, v: Array<F32>, p: Array<F32>, dka: Array<F32>, dva: Array<F32>, r: PO) -> BS:
2488 match r:
2489 case PO{g, k, dq}:
2490 hb4(t, hd, w, hc, k, d, v, p, dq, dva, pacc(Nat.sub(Nat.mul(t, hd), 1n), 0, U32.from_nat(hd), t, U32.from_nat(w), hc, dka, dot0(t, 0, hc, U32.from_nat(t), U32.from_nat(w), g, q)))
2491
2492def hb2(+t: Nat, +hd: Nat, +w: Nat, +kvw: Nat, +hc: U32, +kc: U32, q: Array<F32>, k: Array<F32>, dq: Array<F32>, dka: Array<F32>, dva: Array<F32>, r: GW) -> BS:
2493 match r:
2494 case GW{d, v, p, g}:
2495 hb3(t, hd, w, hc, q, d, v, p, dka, dva, pmul(Nat.sub(Nat.mul(t, hd), 1n), 0, U32.from_nat(hd), t, U32.from_nat(kvw), kc, U32.from_nat(w), hc, dq, dot0(t, 0, kc, 1, U32.from_nat(kvw), g, k)))
2496
2497# P again, then G
2498def hb1(+t: Nat, +hd: Nat, +w: Nat, +kvw: Nat, +sc: F32, +hc: U32, +kc: U32, v: Array<F32>, d: Array<F32>, dq: Array<F32>, dka: Array<F32>, dva: Array<F32>, r: QK) -> BS:
2499 match r:
2500 case QK{q, k, p}:
2501 hb2(t, hd, w, kvw, hc, kc, q, k, dq, dka, dva, grows(Nat.sub(t, 1n), 0n, t, hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, grow(0n, t, hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, d, v, p, zeros_a(Nat.mul(t, t)))))
2502
2503# one head's backward pass at head_cols' columns (no head index, as head_fwd)
2504def head_bwd(+dm: Dm, +cols: HC, s: BS) -> BS:
2505 match dm cols s:
2506 case Dm{+t, +g, +hd, +w, +kvw, +sc} HC{+hc, +kc} BS{q, k, v, d, dq, dka, dva}:
2507 hb1(t, hd, w, kvw, sc, hc, kc, v, d, dq, dka, dva, probs(t, hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, q, k))
2508
2509def bheads(left: Nat, +hi: Nat, +dm: Dm, s: BS) -> BS:
2510 match left:
2511 case 0n:
2512 s
2513 case 1n+n:
2514 bheads(n, 1n+hi, dm, head_bwd(dm, head_cols(dm, hi), s))
2515
2516type BG is Type:
2517 BG{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, dq: Array<F32>, dk: Array<F32>, dv: Array<F32>}
2518
2519# the group's summed dK and dV into key/value head j's columns of dk and dv
2520def bg_fin(+t: Nat, +hd: Nat, +kvw: Nat, +j: Nat, dk: Array<F32>, dv: Array<F32>, s: BS) -> BG:
2521 match s:
2522 case BS{q, k, v, d, dq, dka, dva}:
2523 BG{q, k, v, d, dq, cols_in(t, hd, U32.from_nat(kvw), U32.from_nat(Nat.mul(j, hd)), dka, dk), cols_in(t, hd, U32.from_nat(kvw), U32.from_nat(Nat.mul(j, hd)), dva, dv)}
2524
2525def bgroup(+dm: Dm, +j: Nat, s: BG) -> BG:
2526 match dm s:
2527 case Dm{+t, +g, +hd, +w, +kvw, +sc} BG{q, k, v, d, dq, dk, dv}:
2528 bg_fin(t, hd, kvw, j, dk, dv, bheads(g, Nat.mul(j, g), dm, BS{q, k, v, d, dq, zeros_a(Nat.mul(t, hd)), zeros_a(Nat.mul(t, hd))}))
2529
2530def bgroups(left: Nat, +j: Nat, +dm: Dm, s: BG) -> BG:
2531 match left:
2532 case 0n:
2533 s
2534 case 1n+n:
2535 bgroups(n, 1n+j, dm, bgroup(dm, j, s))
2536
2537# ---- the sequences: each one's rows copied out into arrays of its own, the sequences computed in a tree of parallel
2538# lets (each leaf holds its own arrays, nothing shared), the results written back in order ----
2539
2540type Halves<-A: Type> is Type:
2541 Halves{l: List<&1, A>, r: List<&1, A>}
2542
2543# the two halves back as one list, in order (what the laws attn_par_fwd and attn_par_bwd compare the trees with)
2544def hcat(-A: Type, h: Halves<A>) -> List<&1, A>:
2545 match h:
2546 case Halves{l, r}:
2547 List.append(&1, A, l, r)
2548
2549# xs cut after its first n (a tail loop: a helper in the fork tree must not hold the stack)
2550def cut(-A: Type, n: Nat, xs: List<&1, A>, acc: List<&1, A>) -> Halves<A>:
2551 match n xs:
2552 case 1n+m Con{x, xt}:
2553 cut(A, m, xt, x <> acc)
2554 case _ _:
2555 Halves{List.reverse(&1, A, acc), xs}
2556
2557type FL is Type:
2558 FL{q: Array<F32>, k: Array<F32>, v: Array<F32>, xs: List<&1, FS>}
2559
2560def fl3(+t: Nat, +w: Nat, qq: Array<F32>, kk: Array<F32>, xs: List<&1, FS>, q: Array<F32>, k: Array<F32>, r: A2) -> FL:
2561 match r:
2562 case A2{vv, v}:
2563 FL{qq, kk, vv, FS{q, k, v, zeros_a(Nat.mul(t, w))} <> xs}
2564
2565def fl2(+t: Nat, +w: Nat, +kvw: Nat, +r0: Nat, qq: Array<F32>, vv: Array<F32>, xs: List<&1, FS>, q: Array<F32>, r: A2) -> FL:
2566 match r:
2567 case A2{kk, k}:
2568 fl3(t, w, qq, kk, xs, q, k, rows_out(t, kvw, r0, vv))
2569
2570def fl1(+t: Nat, +w: Nat, +kvw: Nat, +r0: Nat, kk: Array<F32>, vv: Array<F32>, xs: List<&1, FS>, r: A2) -> FL:
2571 match r:
2572 case A2{qq, q}:
2573 fl2(t, w, kvw, r0, qq, vv, xs, q, rows_out(t, kvw, r0, kk))
2574
2575def fslice(+t: Nat, +w: Nat, +kvw: Nat, +bi: Nat, s: FL) -> FL:
2576 match s:
2577 case FL{qq, kk, vv, xs}:
2578 fl1(t, w, kvw, Nat.mul(bi, t), kk, vv, xs, rows_out(t, w, Nat.mul(bi, t), qq))
2579
2580# every sequence's rows of Q, K and V, newest first
2581def fslices(left: Nat, +bi: Nat, +t: Nat, +w: Nat, +kvw: Nat, s: FL) -> FL:
2582 match left:
2583 case 0n:
2584 s
2585 case 1n+n:
2586 fslices(n, 1n+bi, t, w, kvw, fslice(t, w, kvw, bi, s))
2587
2588def fleaf(xs: List<&1, FS>, +h: Nat, +dm: Dm, acc: List<&1, Array<F32>>) -> List<&1, Array<F32>>:
2589 match xs:
2590 case Nil{}:
2591 List.reverse(&1, Array<F32>, acc)
2592 case Con{x, rest}:
2593 fleaf(rest, h, dm, fs_o(fheads(h, 0n, dm, x)) <> acc)
2594
2595# the n sequences of gs (cut in two halves) in a tree of parallel lets of depth d; the outputs come back in order
2596def fpar(d: Nat, gs: Halves<FS>, +n: Nat, +h: Nat, +dm: Dm) -> List<&1, Array<F32>>:
2597 match d gs:
2598 case 0n Halves{l, r}:
2599 a b = fleaf(l, h, dm, Nil{}) fleaf(r, h, dm, Nil{})
2600 List.append(&1, Array<F32>, a, b)
2601 case 1n++e Halves{l, r}:
2602 a b = fpar(e, cut(FS, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm) fpar(e, cut(FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm)
2603 List.append(&1, Array<F32>, a, b)
2604
2605# the blocks (t rows, c wide each) into o from row r0, in order
2606def blocks_in(xs: List<&1, Array<F32>>, +t: Nat, +c: Nat, +r0: Nat, o: Array<F32>) -> Array<F32>:
2607 match xs:
2608 case Nil{}:
2609 o
2610 case Con{x, rest}:
2611 blocks_in(rest, t, c, Nat.add(r0, t), rows_in(t, c, r0, x, o))
2612
2613type BL is Type:
2614 BL{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, xs: List<&1, BG>}
2615
2616def bl4(+t: Nat, +w: Nat, +kvw: Nat, qq: Array<F32>, kk: Array<F32>, vv: Array<F32>, xs: List<&1, BG>, q: Array<F32>, k: Array<F32>, v: Array<F32>, r: A2) -> BL:
2617 match r:
2618 case A2{dd, d}:
2619 BL{qq, kk, vv, dd, BG{q, k, v, d, zeros_a(Nat.mul(t, w)), zeros_a(Nat.mul(t, kvw)), zeros_a(Nat.mul(t, kvw))} <> xs}
2620
2621def bl3(+t: Nat, +w: Nat, +kvw: Nat, +r0: Nat, qq: Array<F32>, kk: Array<F32>, dd: Array<F32>, xs: List<&1, BG>, q: Array<F32>, k: Array<F32>, r: A2) -> BL:
2622 match r:
2623 case A2{vv, v}:
2624 bl4(t, w, kvw, qq, kk, vv, xs, q, k, v, rows_out(t, w, r0, dd))
2625
2626def bl2(+t: Nat, +w: Nat, +kvw: Nat, +r0: Nat, qq: Array<F32>, vv: Array<F32>, dd: Array<F32>, xs: List<&1, BG>, q: Array<F32>, r: A2) -> BL:
2627 match r:
2628 case A2{kk, k}:
2629 bl3(t, w, kvw, r0, qq, kk, dd, xs, q, k, rows_out(t, kvw, r0, vv))
2630
2631def bl1(+t: Nat, +w: Nat, +kvw: Nat, +r0: Nat, kk: Array<F32>, vv: Array<F32>, dd: Array<F32>, xs: List<&1, BG>, r: A2) -> BL:
2632 match r:
2633 case A2{qq, q}:
2634 bl2(t, w, kvw, r0, qq, vv, dd, xs, q, rows_out(t, kvw, r0, kk))
2635
2636def bslice(+t: Nat, +w: Nat, +kvw: Nat, +bi: Nat, s: BL) -> BL:
2637 match s:
2638 case BL{qq, kk, vv, dd, xs}:
2639 bl1(t, w, kvw, Nat.mul(bi, t), kk, vv, dd, xs, rows_out(t, w, Nat.mul(bi, t), qq))
2640
2641def bslices(left: Nat, +bi: Nat, +t: Nat, +w: Nat, +kvw: Nat, s: BL) -> BL:
2642 match left:
2643 case 0n:
2644 s
2645 case 1n+n:
2646 bslices(n, 1n+bi, t, w, kvw, bslice(t, w, kvw, bi, s))
2647
2648type DB is Type:
2649 DB{dq: Array<F32>, dk: Array<F32>, dv: Array<F32>}
2650
2651def db_of(s: BG) -> DB:
2652 match s:
2653 case BG{q, k, v, d, dq, dk, dv}:
2654 DB{dq, dk, dv}
2655
2656def bleaf(xs: List<&1, BG>, +kv: Nat, +dm: Dm, acc: List<&1, DB>) -> List<&1, DB>:
2657 match xs:
2658 case Nil{}:
2659 List.reverse(&1, DB, acc)
2660 case Con{x, rest}:
2661 bleaf(rest, kv, dm, db_of(bgroups(kv, 0n, dm, x)) <> acc)
2662
2663def bpar(d: Nat, gs: Halves<BG>, +n: Nat, +kv: Nat, +dm: Dm) -> List<&1, DB>:
2664 match d gs:
2665 case 0n Halves{l, r}:
2666 a b = bleaf(l, kv, dm, Nil{}) bleaf(r, kv, dm, Nil{})
2667 List.append(&1, DB, a, b)
2668 case 1n++e Halves{l, r}:
2669 a b = bpar(e, cut(BG, TA.half(TA.half(n)), l, Nil{}), TA.half(n), kv, dm) bpar(e, cut(BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), kv, dm)
2670 List.append(&1, DB, a, b)
2671
2672type D3 is Type:
2673 D3{dq: Array<F32>, dk: Array<F32>, dv: Array<F32>}
2674
2675def dblocks_in(xs: List<&1, DB>, +t: Nat, +w: Nat, +kvw: Nat, +r0: Nat, o: D3) -> D3:
2676 match xs o:
2677 case Nil{} o:
2678 o
2679 case Con{DB{bq, bk, bv}, rest} D3{dq, dk, dv}:
2680 dblocks_in(rest, t, w, kvw, Nat.add(r0, t), D3{rows_in(t, w, r0, bq, dq), rows_in(t, kvw, r0, bk, dk), rows_in(t, kvw, r0, bv, dv)})
2681
2682# ---- the layer ----
2683
2684type AttnDims is Data:
2685 AttnDims{b: Nat, t: Nat, kv: Nat, g: Nat, hd: Nat, sc: F32}
2686
2687def dm_of(+w: Nat, +kvw: Nat, +dm: AttnDims) -> Dm:
2688 match dm:
2689 case AttnDims{+b, +t, +kv, +g, +hd, +sc}:
2690 Dm{t, g, hd, w, kvw, sc}
2691
2692# whether the sequences run in the tree: from 8 sequences on, or from 2 when each sequence is big. With the tiny Llama's
2693# sequences (t = 16, w = 32) the tree cost more than it gave at 2 sequences (0.48 against 0.37 ms on 16 threads) and
2694# gained from 8 on (1.32 against 1.52 ms; 1.5x at 32, 1.7x at 128); with SmolLM2's (t = 128, w = 576, about half a
2695# second of a step each) the two sequences one after another kept a single thread busy for a third of the step (NOTES,
2696# 2026-10-21). t w counts in U32 (a Nat comparison walks its numbers one unit at a time). The answer never changes the
2697# result: the laws attn_par_fwd and attn_par_bwd show the tree gives the sequences in order, as one after another.
2698def attn_tree(+b: Nat, +t: Nat, +w: Nat) -> Bool:
2699 Bool.or(Nat.is_le(8n, b), Bool.and(Nat.is_le(2n, b), U32.is_le(4096, (U32.from_nat(t) * U32.from_nat(w) : U32))))
2700
2701def fgo(big: Bool, +b: Nat, +h: Nat, +dm: Dm, xs: List<&1, FS>) -> List<&1, Array<F32>>:
2702 match big:
2703 case True{}:
2704 fpar(Nat.sub(depth(b), 1n), cut(FS, TA.half(b), xs, Nil{}), b, h, dm)
2705 case False{}:
2706 fleaf(xs, h, dm, Nil{})
2707
2708def bgo(big: Bool, +b: Nat, +kv: Nat, +dm: Dm, xs: List<&1, BG>) -> List<&1, DB>:
2709 match big:
2710 case True{}:
2711 bpar(Nat.sub(depth(b), 1n), cut(BG, TA.half(b), xs, Nil{}), b, kv, dm)
2712 case False{}:
2713 bleaf(xs, kv, dm, Nil{})
2714
2715def af2(-n: Nat, +w: Nat, +b: Nat, +h: Nat, +dm: Dm, s: FL) -> Fwd<Unit, TA.Mat<n, w>, AttnS>:
2716 match dm s:
2717 case Dm{+t, +g, +hd, +ww, +kvw, +sc} FL{q, k, v, xs}:
2718 Fwd{Unit{}, TA.Mat{blocks_in(fgo(attn_tree(b, t, w), b, h, dm, List.reverse(&1, FS, xs)), t, w, 0n, zeros_a(Nat.mul(Nat.mul(b, t), w)))}, AttnS{q, k, v}}
2719
2720def attn_fwd(+n: Nat, +w: Nat, +kvw: Nat, +dm: AttnDims, x: Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>) -> Fwd<Unit, TA.Mat<n, w>, AttnS>:
2721 match dm x:
2722 case AttnDims{+b, +t, +kv, +g, +hd, +sc} Two{TA.Mat{q}, Two{TA.Mat{k}, TA.Mat{v}}}:
2723 af2(n, w, b, Nat.mul(kv, g), dm_of(w, kvw, dm), fslices(b, 0n, t, w, kvw, FL{q, k, v, Nil{}}))
2724
2725def ab3(-n: Nat, -w: Nat, -kvw: Nat, o: D3) -> Bwd<Unit, Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
2726 match o:
2727 case D3{dq, dk, dv}:
2728 Bwd{Unit{}, Two{TA.Mat{dq}, Two{TA.Mat{dk}, TA.Mat{dv}}}, Unit{}}
2729
2730def ab2(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +kv: Nat, +dm: Dm, s: BL) -> Bwd<Unit, Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
2731 match dm s:
2732 case Dm{+t, +g, +hd, +ww, +kw, +sc} BL{q, k, v, d, xs}:
2733 ab3(n, w, kvw, dblocks_in(bgo(attn_tree(b, t, w), b, kv, dm, List.reverse(&1, BG, xs)), t, w, kvw, 0n, D3{zeros_a(Nat.mul(n, w)), zeros_a(Nat.mul(n, kvw)), zeros_a(Nat.mul(n, kvw))}))
2734
2735def attn_bwd(+n: Nat, +w: Nat, +kvw: Nat, +dm: AttnDims, s: AttnS, dy: TA.Mat<n, w>) -> Bwd<Unit, Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
2736 match dm s dy:
2737 case AttnDims{+b, +t, +kv, +g, +hd, +sc} AttnS{q, k, v} TA.Mat{d}:
2738 ab2(n, w, kvw, b, kv, dm_of(w, kvw, dm), bslices(b, 0n, t, w, kvw, BL{q, k, v, d, Nil{}}))
2739
2740# the core of attention, after the projections and RoPE; the erased arguments are the shape facts it relies on
2741def Attention(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, -e1: {n == Nat.mul(b, t) : Nat}, -e2: {w == Nat.mul(Nat.mul(kv, g), hd) : Nat}, -e3: {kvw == Nat.mul(kv, hd) : Nat}) -> Layer<Unit, Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>, TA.Mat<n, w>, AttnS>:
2742 Layer{p => x => attn_fwd(n, w, kvw, AttnDims{b, t, kv, g, hd, sc}, x), p => s => dy => g2 => attn_bwd(n, w, kvw, AttnDims{b, t, kv, g, hd, sc}, s, dy), op => p => g2 => st => no_step(Unit, op, p, g2, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0}
2743
2744# =====================================================================================================
2745# Losses over arrays (cross-entropy, MSE, NLL)
2746# =====================================================================================================
2747
2748# ---- cross-entropy of the logits z (n x c) against one label per row, mean over the rows (PyTorch's F.cross_entropy,
2749# reduction="mean"); its gradient is (softmax(z) - onehot) / n. Over z's own array, a row at a time, with
2750# bend-ml-tensor-array's Mat.softmax_ce's operations in its order, so its bits. ----
2751
2752def hit(b: Bool) -> F32:
2753 match b:
2754 case True{}:
2755 1.0
2756 case False{}:
2757 0.0
2758
2759# the sum of a[0] .. a[m-1] as a balanced tree of additions (pairwise summation), in place: each pass adds neighbours
2760# (a[i] = a[2i] + a[2i+1], an odd last one carried), until one number is left. Its rounding error grows with log2(m),
2761# not m, as in PyTorch's cascade sum: a running sum of 295k squares was 1e-4 off (NOTES, 2026-10-17)
2762def ts_add(+i: U32, +x: F32, r: Array<F32> & F32) -> Array<F32>:
2763 match r:
2764 case Tuple{a, +y}:
2765 Array.set(F32, a, i, (x + y : F32))
2766
2767def ts_pair(+i: U32, r: Array<F32> & F32) -> Array<F32>:
2768 match r:
2769 case Tuple{a, +x}:
2770 ts_add(i, x, Array.get(F32, a, ((i + i : U32) + 1 : U32)))
2771
2772def ts_pairs(k: Nat, +i: U32, a: Array<F32>) -> Array<F32>:
2773 match k:
2774 case 0n:
2775 a
2776 case 1n+q:
2777 ts_pairs(q, (i + 1 : U32), ts_pair(i, Array.get(F32, a, (i + i : U32))))
2778
2779def ts_put(+i: U32, r: Array<F32> & F32) -> Array<F32>:
2780 match r:
2781 case Tuple{a, +x}:
2782 Array.set(F32, a, i, x)
2783
2784def ts_carry(odd: Bool, +m: U32, a: Array<F32>) -> Array<F32>:
2785 match odd:
2786 case True{}:
2787 ts_put(U32.div(m, 2), Array.get(F32, a, (m - 1 : U32)))
2788 case False{}:
2789 a
2790
2791def ts_pass(+m: U32, a: Array<F32>) -> Array<F32>:
2792 ts_carry(U32.is_eq(U32.mod(m, 2), 1), m, ts_pairs(U32.to_nat(U32.div(m, 2)), 0, a))
2793
2794def ts_first(r: Array<F32> & F32) -> F32:
2795 match r:
2796 case Tuple{a, x}:
2797 x
2798
2799def ts_go(f: Nat, done: Bool, +m: U32, a: Array<F32>) -> F32:
2800 match f done:
2801 case _ True{}:
2802 ts_first(Array.get(F32, a, 0))
2803 case 0n False{}:
2804 ts_first(Array.get(F32, a, 0))
2805 case 1n+g False{}:
2806 ts_go(g, U32.is_le(U32.div((m + 1 : U32), 2), 1), U32.div((m + 1 : U32), 2), ts_pass(m, a))
2807
2808def tree_sum(+m: U32, a: Array<F32>) -> F32:
2809 ts_go(32n, U32.is_le(m, 1), m, a)
2810
2811# Cross-entropy's softmax over a row of c classes: as psm, but the row's exponentials are summed as a balanced tree
2812# (tree_sum, on a copy), since a running F32 sum of SmolLM2's 49152 classes was 3.4 tolerances off PyTorch's loss
2813# (NOTES, M10). psm itself keeps its order: attention's rows are short and the law causal reads its count.
2814# a row's probabilities, its max m and log of its sum of exp(x - m) (log_softmax's pieces)
2815type CSM is Type:
2816 CSM{a: Array<F32>, m: F32, ls: F32}
2817
2818def cesm4(+k: Nat, +off: U32, +m: F32, +s: F32, a: Array<F32>) -> CSM:
2819 CSM{pc_a(pdiv(k, off, s, 0n, Array.get(F32, a, off))), m, F32.log(s)}
2820
2821def cesm3(+k: Nat, +off: U32, +m: F32, r: A2) -> CSM:
2822 match r:
2823 case A2{a, sc}:
2824 cesm4(k, off, m, tree_sum(U32.from_nat(1n+k), sc), a)
2825
2826def cesm_e(+k: Nat, +off: U32, +m: F32, r: AS) -> CSM:
2827 match r:
2828 case AS{a, s}:
2829 cesm3(k, off, m, copy_n(k, off, 0, Array.new(F32, depth(1n+k), 0.0), Array.get(F32, a, off)))
2830
2831def cesm2(+k: Nat, +off: U32, r: AS) -> CSM:
2832 match r:
2833 case AS{a, +m}:
2834 cesm_e(k, off, m, pexp(k, off, m, 0.0, Array.get(F32, a, off)))
2835
2836def ce_sm(+c: Nat, +off: U32, z: Array<F32>) -> CSM:
2837 cesm2(Nat.sub(c, 1n), off, pmax0(Nat.sub(c, 1n), off, Array.get(F32, z, off)))
2838
2839# z[idx] = (p - [j == y]) / nf for j = jj..jj + k, p the row's probabilities
2840def ce_g(k: Nat, +idx: U32, +j: U32, +y: U32, +nf: F32, r: Array<F32> & F32) -> Array<F32>:
2841 match k r:
2842 case 0n Tuple{z, +x}:
2843 Array.set(F32, z, idx, ((x - hit(U32.is_eq(j, y)) : F32) / nf : F32))
2844 case 1n+q Tuple{z, +x}:
2845 ce_g(q, (idx + 1 : U32), (j + 1 : U32), y, nf, Array.get(F32, Array.set(F32, z, idx, ((x - hit(U32.is_eq(j, y)) : F32) / nf : F32)), (idx + 1 : U32)))
2846
2847# loss - log_softmax(x)[y], log_softmax as PyTorch computes it: (x[y] - m) - log(sum exp(x - m)) (since 2026-10-18; it
2848# was loss - log(max(p[y], 1e-12)), which capped a row's loss at 27.63 where PyTorch's grows, NOTES)
2849def ce_pick2(+xy: F32, loss: F32, r: CSM) -> AS:
2850 match r:
2851 case CSM{z, +m, +ls}:
2852 AS{z, (loss - ((xy - m : F32) - ls : F32) : F32)}
2853
2854def ce_pick3(+c: Nat, +off: U32, loss: F32, r: Array<F32> & F32) -> AS:
2855 match r:
2856 case Tuple{z, +xy}:
2857 ce_pick2(xy, loss, ce_sm(c, off, z))
2858
2859def ce_bad(loss: F32, r: CSM) -> AS:
2860 match r:
2861 case CSM{z, m, ls}:
2862 AS{z, (loss - F32.log(F32.max(0.0, 0.000000000001)) : F32)}
2863
2864# a label past the last class picks probability 0.0, clamped at 1e-12, as tensor-array's pick
2865def ce_pick(ok: Bool, +c: Nat, +off: U32, +idx: U32, loss: F32, z: Array<F32>) -> AS:
2866 match ok:
2867 case True{}:
2868 ce_pick3(c, off, loss, Array.get(F32, z, idx))
2869 case False{}:
2870 ce_bad(loss, ce_sm(c, off, z))
2871
2872def ce_row2(+c: Nat, +off: U32, +y: Nat, +nf: F32, r: AS) -> AS:
2873 match r:
2874 case AS{z, +loss}:
2875 AS{ce_g(Nat.sub(c, 1n), off, 0, U32.from_nat(y), nf, Array.get(F32, z, off)), loss}
2876
2877def ce_row(+c: Nat, +off: U32, +y: Nat, +nf: F32, s: AS) -> AS:
2878 match s:
2879 case AS{z, loss}:
2880 ce_row2(c, off, y, nf, ce_pick(Nat.is_le(1n+y, c), c, off, (off + U32.from_nat(y) : U32), loss, z))
2881
2882def ce_rows(labels: List<&2, Nat>, +c: Nat, +i: U32, +nf: F32, s: AS) -> AS:
2883 match labels:
2884 case Nil{}:
2885 s
2886 case Con{+y, t}:
2887 ce_rows(t, c, (i + 1 : U32), nf, ce_row(c, (i * U32.from_nat(c) : U32), y, nf, s))
2888
2889def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: AS) -> LossOut<TA.Mat<n, c>>:
2890 match r:
2891 case AS{d, +l}:
2892 LossOut{(l / nf : F32), TA.Mat{d}}
2893
2894def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat<n, c>, labels: List<&2, Nat>) -> LossOut<TA.Mat<n, c>>:
2895 match z:
2896 case TA.Mat{d}:
2897 ce_fin(n, c, F32.from_nat(n), ce_rows(labels, c, 0, F32.from_nat(n), AS{d, 0.0}))
2898
2899
2900type D3S is Type:
2901 D3S{y: Array<F32>, t: Array<F32>, q: Array<F32>}
2902
2903# over every number from idx on (k more): d = y - t, q[idx] = d·d, y[idx] = d * kk
2904def mse_l(k: Nat, +idx: U32, +kk: F32, q: Array<F32>, ry: Array<F32> & F32, rt: Array<F32> & F32) -> D3S:
2905 match k ry rt:
2906 case 0n Tuple{y, +a} Tuple{t, +b}:
2907 D3S{Array.set(F32, y, idx, ((a - b : F32) * kk : F32)), t, Array.set(F32, q, idx, ((a - b : F32) * (a - b : F32) : F32))}
2908 case 1n+n Tuple{y, +a} Tuple{t, +b}:
2909 mse_l(n, (idx + 1 : U32), kk, Array.set(F32, q, idx, ((a - b : F32) * (a - b : F32) : F32)), Array.get(F32, Array.set(F32, y, idx, ((a - b : F32) * kk : F32)), (idx + 1 : U32)), Array.get(F32, t, (idx + 1 : U32)))
2910
2911def mse_fin(-n: Nat, -c: Nat, +m: U32, +nc: F32, r: D3S) -> LossOut<TA.Mat<n, c>>:
2912 match r:
2913 case D3S{y, t, q}:
2914 LossOut{(tree_sum(m, q) / nc : F32), TA.Mat{y}}
2915
2916# mean squared error over every number (PyTorch's F.mse_loss); its gradient is 2 (y - t) / (n*c); the squares summed
2917# as a tree
2918def MSE(+n: Nat, +c: Nat, y: TA.Mat<n, c>, t: TA.Mat<n, c>) -> LossOut<TA.Mat<n, c>>:
2919 match y t:
2920 case TA.Mat{ya} TA.Mat{ta}:
2921 mse_fin(n, c, U32.from_nat(Nat.mul(n, c)), F32.from_nat(Nat.mul(n, c)), mse_l(Nat.sub(Nat.mul(n, c), 1n), 0, (2.0 / F32.from_nat(Nat.mul(n, c)) : F32), zeros_a(Nat.mul(n, c)), Array.get(F32, ya, 0), Array.get(F32, ta, 0)))
2922
2923type NS is Type:
2924 NS{lp: Array<F32>, g: Array<F32>, s: F32}
2925
2926def nll_hit(+idx: U32, +v: F32, g: Array<F32>, acc: F32, r: Array<F32> & F32) -> NS:
2927 match r:
2928 case Tuple{lp, +x}:
2929 NS{lp, Array.set(F32, g, idx, v), (acc + x : F32)}
2930
2931# a label past the last class picks 0.0 and leaves its gradient row zero, as l_at and onehot do
2932def nll_one(ok: Bool, +idx: U32, +v: F32, acc: F32, lp: Array<F32>, g: Array<F32>) -> NS:
2933 match ok:
2934 case True{}:
2935 nll_hit(idx, v, g, acc, Array.get(F32, lp, idx))
2936 case False{}:
2937 NS{lp, g, (acc + 0.0 : F32)}
2938
2939# the sum of the picked log-probabilities (in row order, from +0.0), and the gradient (-1/n at each label)
2940def nll_rows(left: Nat, labels: List<&2, Nat>, +i: U32, +c: Nat, +v: F32, st: NS) -> NS:
2941 match left labels st:
2942 case 1n+q Con{+y, t} NS{lp, g, acc}:
2943 nll_rows(q, t, (i + 1 : U32), c, v, nll_one(Nat.is_le(1n+y, c), ((i * U32.from_nat(c) : U32) + U32.from_nat(y) : U32), v, acc, lp, g))
2944 case _ _ st:
2945 st
2946
2947def nll_fin(-c: Nat, +n: Nat, r: NS) -> LossOut<TA.Mat<n, c>>:
2948 match r:
2949 case NS{lp, g, +s}:
2950 LossOut{F32.neg((s / F32.from_nat(n) : F32)), TA.Mat{g}}
2951
2952# negative log-likelihood of log-probabilities (PyTorch's F.nll_loss, mean over the rows)
2953def NLL(+n: Nat, +c: Nat, logp: TA.Mat<n, c>, labels: List<&2, Nat>) -> LossOut<TA.Mat<n, c>>:
2954 match logp:
2955 case TA.Mat{lp}:
2956 nll_fin(c, n, nll_rows(n, labels, 0, c, F32.neg((1.0 / F32.from_nat(n) : F32)), NS{lp, zeros_a(Nat.mul(n, c)), 0.0}))
2957
2958# =====================================================================================================
2959# Checkpoints: parameters, only the buffers in use, and Adam powers, as little-endian F32 (the gradient is zero
2960# between steps and is not saved). A run stopped after step k and resumed from its checkpoint computes the same bits
2961# as a run that never stopped.
2962# =====================================================================================================
2963
2964def byte_of(+u: U32, +k: Nat) -> U32:
2965 U32.and(U32.shrn(u, k), 255)
2966
2967def read_all(+path: String) -> IO(List<&2, U32>):
2968 do IO<List<&2, U32>>:
2969 f : File <- IO.try(File, File.open(path, "r"))
2970 m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_bytes(f, 2000000000)
2971 rd_done(m)
2972
2973# ---- packed BF16 moments in a checkpoint (CAdamWBF16) ----
2974# A PTree is written as F32 cells holding U32 bits, before 1 - b1^t and 1 - b2^t, in the order the step walks it:
2975# PNew: 0, v16 (1 or 0); PNode: 1, its left part, its right part; PLeaf: 2, n, m's ceil(n / 2) words, v's; PLeafV: 3, n,
2976# m's words, v's n F32 numbers. The tree carries its own shape and sizes, so a fresh state (all PNew) reads it back.
2977
2978def ph_hw(+n: U32) -> U32:
2979 U32.shrn((n + 1 : U32), 1n)
2980
2981def ph_bool(b: Bool) -> U32:
2982 match b:
2983 case True{}: 1
2984 case False{}: 0
2985
2986type PTU is Type:
2987 PTU{t: PTree, n: U32}
2988
2989def ptu_node(ra: PTU, rb: PTU) -> PTU:
2990 match ra rb:
2991 case PTU{a, na} PTU{b, nb}: PTU{PNode{a, b}, ((1 + na : U32) + nb : U32)}
2992
2993# the number of cells a tree takes (the tree comes back)
2994def pt_cells(t: PTree) -> PTU:
2995 match t:
2996 case PNew{v16}: PTU{PNew{v16}, 2}
2997 case PLeaf{+n, m, v}: PTU{PLeaf{n, m, v}, (2 + (2 * ph_hw(n) : U32) : U32)}
2998 case PLeafV{+n, m, v}: PTU{PLeafV{n, m, v}, ((2 + ph_hw(n) : U32) + n : U32)}
2999 case PNode{a, b}: ptu_node(pt_cells(a), pt_cells(b))
3000
3001type PTA is Type:
3002 PTA{t: PTree, xs: Array<F32>, off: U32}
3003
3004type UA is Type:
3005 UA{a: Array<U32>, xs: Array<F32>}
3006
3007# words i.. of a packed buffer into xs from off, as F32 cells with their bits (both arrays come back)
3008def words_out(left: Nat, +i: U32, +off: U32, ra: Array<U32> & U32, xs: Array<F32>) -> UA:
3009 match left ra:
3010 case 0n Tuple{a, w}: UA{a, xs}
3011 case 1n+q Tuple{a, +w}:
3012 words_out(q, (i + 1 : U32), (off + 1 : U32), Array.get(U32, a, (i + 1 : U32)), Array.set(F32, xs, off, ph_u2f(w)))
3013
3014def words_in(left: Nat, +i: U32, +off: U32, rx: Array<F32> & F32, a: Array<U32>) -> UA:
3015 match left rx:
3016 case 0n Tuple{xs, x}: UA{a, xs}
3017 case 1n+q Tuple{xs, x}:
3018 words_in(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, xs, (off + 1 : U32)), Array.set(U32, a, i, ph_f2u(x)))
3019
3020def pta_then(r: PTA, k: PTree -> Array<F32> -> U32 -> PTA) -> PTA:
3021 match r:
3022 case PTA{t, xs, off}: k(t, xs, off)
3023
3024def pta_node(a: PTree, r: PTA) -> PTA:
3025 match r:
3026 case PTA{b, xs, off}: PTA{PNode{a, b}, xs, off}
3027
3028def pts_v(+n: U32, m: Array<U32>, +off: U32, r: UA) -> PTA:
3029 match r:
3030 case UA{v, xs}: PTA{PLeaf{n, m, v}, xs, (off + ph_hw(n) : U32)}
3031
3032def pts_m(+n: U32, v: Array<U32>, +off: U32, r: UA) -> PTA:
3033 match r:
3034 case UA{m, xs}: pts_v(n, m, (off + ph_hw(n) : U32), words_out(U32.to_nat(ph_hw(n)), 0, (off + ph_hw(n) : U32), Array.get(U32, v, 0), xs))
3035
3036def ptsv_v(+n: U32, m: Array<U32>, +off: U32, r: TA.P2) -> PTA:
3037 match r:
3038 case TA.P2{v, xs}: PTA{PLeafV{n, m, v}, xs, (off + n : U32)}
3039
3040def ptsv_m(+n: U32, v: Array<F32>, +off: U32, r: UA) -> PTA:
3041 match r:
3042 case UA{m, xs}: ptsv_v(n, m, (off + ph_hw(n) : U32), store_cells(U32.to_nat(n), 0, (off + ph_hw(n) : U32), Array.get(F32, v, 0), xs))
3043
3044# the tree's cells into xs from off
3045def pt_store(t: PTree, xs: Array<F32>, +off: U32) -> PTA:
3046 match t:
3047 case PNew{+v16}:
3048 PTA{PNew{v16}, Array.set(F32, Array.set(F32, xs, off, ph_u2f(0)), (off + 1 : U32), ph_u2f(ph_bool(v16))), (off + 2 : U32)}
3049 case PLeaf{+n, m, v}:
3050 pts_m(n, v, (off + 2 : U32), words_out(U32.to_nat(ph_hw(n)), 0, (off + 2 : U32), Array.get(U32, m, 0), Array.set(F32, Array.set(F32, xs, off, ph_u2f(2)), (off + 1 : U32), ph_u2f(n))))
3051 case PLeafV{+n, m, v}:
3052 ptsv_m(n, v, (off + 2 : U32), words_out(U32.to_nat(ph_hw(n)), 0, (off + 2 : U32), Array.get(U32, m, 0), Array.set(F32, Array.set(F32, xs, off, ph_u2f(3)), (off + 1 : U32), ph_u2f(n))))
3053 case PNode{a, b}:
3054 pta_then(pt_store(a, Array.set(F32, xs, off, ph_u2f(1)), (off + 1 : U32)), ta => xs2 => off2 => pta_node(ta, pt_store(b, xs2, off2)))
3055
3056# reading a leaf back: fresh buffers filled from xs
3057def ptl_v(+n: U32, m: Array<U32>, +off: U32, r: UA) -> PTA:
3058 match r:
3059 case UA{v, xs}: PTA{PLeaf{n, m, v}, xs, (off + ph_hw(n) : U32)}
3060
3061def ptl_m(+n: U32, +off: U32, r: UA) -> PTA:
3062 match r:
3063 case UA{m, xs}: ptl_v(n, m, (off + ph_hw(n) : U32), words_in(U32.to_nat(ph_hw(n)), 0, (off + ph_hw(n) : U32), Array.get(F32, xs, (off + ph_hw(n) : U32)), ph_words(U32.to_nat(n))))
3064
3065def ptlv_v(+n: U32, m: Array<U32>, +off: U32, r: TA.P2) -> PTA:
3066 match r:
3067 case TA.P2{xs, v}: PTA{PLeafV{n, m, v}, xs, (off + n : U32)}
3068
3069def ptlv_m(+n: U32, +off: U32, r: UA) -> PTA:
3070 match r:
3071 case UA{m, xs}: ptlv_v(n, m, (off + ph_hw(n) : U32), load_cells(U32.to_nat(n), 0, (off + ph_hw(n) : U32), Array.get(F32, xs, (off + ph_hw(n) : U32)), Array.new(F32, depth(U32.to_nat(n)), 0.0)))
3072
3073# what a tag says
3074type PK is Data:
3075 KNew{}
3076 KNode{}
3077 KLeaf{}
3078 KLeafV{}
3079
3080def pt_kind2(z: Bool, o: Bool, l: Bool) -> PK:
3081 match z o l:
3082 case True{} _ _: KNew{}
3083 case False{} True{} _: KNode{}
3084 case False{} False{} True{}: KLeaf{}
3085 case False{} False{} False{}: KLeafV{}
3086
3087def pt_kind(+tag: U32) -> PK:
3088 pt_kind2(U32.is_eq(tag, 0), U32.is_eq(tag, 1), U32.is_eq(tag, 2))
3089
3090# the tag at off and the cell after it, to a continuation
3091def pt_next(+off: U32, r: Array<F32> & F32, k: PK -> (Array<F32> & F32) -> PTA) -> PTA:
3092 match r:
3093 case Tuple{xs, x}: k(pt_kind(ph_f2u(x)), Array.get(F32, xs, (off + 1 : U32)))
3094
3095# the part at off, of kind k, with its second cell rn already read; fuel bounds the depth (a model's pairs nest far
3096# less than 1000 deep)
3097def pt_read(fuel: Nat, k: PK, +off: U32, rn: Array<F32> & F32) -> PTA:
3098 match fuel k rn:
3099 case 0n _ Tuple{xs, x}:
3100 PTA{PNew{False{}}, xs, off}
3101 case 1n+f KNew{} Tuple{xs, x}:
3102 PTA{PNew{U32.is_eq(ph_f2u(x), 1)}, xs, (off + 2 : U32)}
3103 case 1n++f KNode{} Tuple{xs, x}:
3104 pta_then(pt_read(f, pt_kind(ph_f2u(x)), (off + 1 : U32), Array.get(F32, xs, (off + 2 : U32))), ta => xs2 => +off2 => pta_node(ta, pt_next(off2, Array.get(F32, xs2, off2), kb => rb => pt_read(f, kb, off2, rb))))
3105 case 1n+f KLeaf{} Tuple{xs, x}:
3106 +n = ph_f2u(x)
3107 ptl_m(n, (off + 2 : U32), words_in(U32.to_nat(ph_hw(n)), 0, (off + 2 : U32), Array.get(F32, xs, (off + 2 : U32)), ph_words(U32.to_nat(n))))
3108 case 1n+f KLeafV{} Tuple{xs, x}:
3109 +n = ph_f2u(x)
3110 ptlv_m(n, (off + 2 : U32), words_in(U32.to_nat(ph_hw(n)), 0, (off + 2 : U32), Array.get(F32, xs, (off + 2 : U32)), ph_words(U32.to_nat(n))))
3111
3112# a tree from xs at off
3113def pt_load(xs: Array<F32>, +off: U32) -> PTA:
3114 pt_next(off, Array.get(F32, xs, off), k => rn => pt_read(1000n, k, off, rn))
3115
3116type TrainPL<-P: Type> is Type:
3117 TrainPL{tr: Train<P>, xs: List<&2, F32>}
3118
3119def td0(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>) -> TrainPL<P>:
3120 match rp:
3121 case PL{p, xp}: TrainPL{Train{p, g, NoMom{}, p1, p2, l}, xp}
3122
3123def td1(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>, rm: PL<P>) -> TrainPL<P>:
3124 match rp rm:
3125 case PL{p, xp} PL{mm, xm}: TrainPL{Train{p, g, OneMom{mm}, p1, p2, l}, List.append(&2, F32, xp, xm)}
3126
3127def td3(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>, rm: PL<P>, rv: PL<P>) -> TrainPL<P>:
3128 match rp rm rv:
3129 case PL{p, xp} PL{m, xm} PL{v, xv}:
3130 TrainPL{Train{p, g, TwoMom{m, v}, p1, p2, l}, List.append(&2, F32, xp, List.append(&2, F32, xm, List.append(&2, F32, xv, [p1, p2])))}
3131
3132# the first k cells of an array, in order
3133def cells_l(k: Nat, +i: U32, acc: List<&2, F32>, ra: Array<F32> & F32) -> List<&2, F32>:
3134 match k ra:
3135 case 0n Tuple{a, v}:
3136 List.reverse(&2, F32, acc)
3137 case 1n+q Tuple{a, +v}:
3138 cells_l(q, (i + 1 : U32), v <> acc, Array.get(F32, a, (i + 1 : U32)))
3139
3140def tdp2(-P: Type, p: P, g: P, +p1: F32, +p2: F32, +l: F32, xp: List<&2, F32>, r: PTA) -> TrainPL<P>:
3141 match r:
3142 case PTA{t, xs, +off}:
3143 TrainPL{Train{p, g, Packed{t}, p1, p2, l}, List.append(&2, F32, xp, List.append(&2, F32, cells_l(U32.to_nat(off), 0, [], Array.get(F32, xs, 0)), [p1, p2]))}
3144
3145def tdp(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>, rt: PTU) -> TrainPL<P>:
3146 match rp rt:
3147 case PL{p, xp} PTU{t, +n}:
3148 tdp2(P, p, g, p1, p2, l, xp, pt_store(t, Array.new(F32, depth(U32.to_nat(n)), 0.0), 0))
3149
3150# every number of a training state: the parameters, then the first moments, the second moments, b1^t and b2^t (packed
3151# moments: their tree's cells, pt_store's layout, in place of the two moments)
3152def Train.dump(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, tr: Train<P>) -> TrainPL<P>:
3153 match tr:
3154 case Train{p, g, NoMom{}, +p1, +p2, +l}:
3155 td0(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p))
3156 case Train{p, g, OneMom{mm}, +p1, +p2, +l}:
3157 td1(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p), Layer.dump(P, X, Y, S, m, mm))
3158 case Train{p, g, TwoMom{mm, vv}, +p1, +p2, +l}:
3159 td3(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p), Layer.dump(P, X, Y, S, m, mm), Layer.dump(P, X, Y, S, m, vv))
3160 case Train{p, g, Packed{t}, +p1, +p2, +l}:
3161 tdp(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p), pt_cells(t))
3162
3163def f32_at(xs: List<&2, F32>, n: Nat) -> F32:
3164 match xs n:
3165 case Con{x, t} 0n:
3166 x
3167 case Con{x, t} 1n+q:
3168 f32_at(t, q)
3169 case Nil{} _:
3170 0.0
3171
3172def tr0(-P: Type, g: P, r: PA<P>) -> Train<P>:
3173 match r:
3174 case PA{p, xs, off}: Train{p, g, NoMom{}, 0.0, 0.0, 0.0}
3175
3176def tr_m2(-P: Type, p: P, g: P, r: PA<P>) -> Train<P>:
3177 match r:
3178 case PA{mm, xs, off}: Train{p, g, OneMom{mm}, 0.0, 0.0, 0.0}
3179
3180def tr_m1(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, mm: P, g: P, r: PA<P>) -> Train<P>:
3181 match r:
3182 case PA{p, xs, off}: tr_m2(P, p, g, Layer.load(P, X, Y, S, m, mm, xs, off))
3183
3184def tr_adam2(-P: Type, p: P, mm: P, vv: P, g: P, p1: F32, r: Array<F32> & F32) -> Train<P>:
3185 match r:
3186 case Tuple{xs, p2}: Train{p, g, TwoMom{mm, vv}, p1, p2, 0.0}
3187
3188def tr_adam(-P: Type, p: P, mm: P, vv: P, g: P, r: Array<F32> & F32, +off: U32) -> Train<P>:
3189 match r:
3190 case Tuple{xs, p1}: tr_adam2(P, p, mm, vv, g, p1, Array.get(F32, xs, off))
3191
3192def tr3(-P: Type, p: P, mm: P, g: P, r: PA<P>) -> Train<P>:
3193 match r:
3194 case PA{vv, xs, +off}: tr_adam(P, p, mm, vv, g, Array.get(F32, xs, off), (off + 1 : U32))
3195
3196def tr2(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, p: P, vv: P, g: P, r: PA<P>) -> Train<P>:
3197 match r:
3198 case PA{mm, xs, off}:
3199 tr3(P, p, mm, g, Layer.load(P, X, Y, S, m, vv, xs, off))
3200
3201def tr1(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, mm: P, vv: P, g: P, r: PA<P>) -> Train<P>:
3202 match r:
3203 case PA{p, xs, off}:
3204 tr2(~P, ~X, ~Y, ~S, ~m, p, vv, g, Layer.load(P, X, Y, S, m, mm, xs, off))
3205
3206def trp4(-P: Type, p: P, g: P, t: PTree, +p1: F32, r: Array<F32> & F32) -> Train<P>:
3207 match r:
3208 case Tuple{xs, p2}: Train{p, g, Packed{t}, p1, p2, 0.0}
3209
3210def trp3(-P: Type, p: P, g: P, t: PTree, r: Array<F32> & F32, +off: U32) -> Train<P>:
3211 match r:
3212 case Tuple{xs, p1}: trp4(P, p, g, t, p1, Array.get(F32, xs, off))
3213
3214def trp2(-P: Type, p: P, g: P, r: PTA) -> Train<P>:
3215 match r:
3216 case PTA{t, xs, +off}: trp3(P, p, g, t, Array.get(F32, xs, off), (off + 1 : U32))
3217
3218def trp(-P: Type, g: P, r: PA<P>) -> Train<P>:
3219 match r:
3220 case PA{p, xs, +off}: trp2(P, p, g, pt_load(xs, off))
3221
3222# a training state filled from the numbers Train.dump wrote (its buffers are overwritten in place; packed moments are
3223# built again from the tree's cells)
3224def Train.restore(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, tr: Train<P>, xs: Array<F32>) -> Train<P>:
3225 match tr:
3226 case Train{p, g, NoMom{}, p1, p2, l}:
3227 tr0(P, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3228 case Train{p, g, OneMom{mm}, p1, p2, l}:
3229 tr_m1(~P, ~X, ~Y, ~S, ~m, mm, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3230 case Train{p, g, TwoMom{mm, vv}, p1, p2, l}:
3231 tr1(~P, ~X, ~Y, ~S, ~m, mm, vv, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3232 case Train{p, g, Packed{t}, p1, p2, l}:
3233 trp(P, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3234
3235# ---- writing numbers straight from an array, a chunk at a time ----
3236# File.write_bytes takes a list with one cell per byte; one list for a whole model would take several GB at SmolLM2's
3237# sizes, so each chunk of WCHUNK numbers gets its own list, written and dropped before the next.
3238
3239def WCHUNK() -> U32:
3240 65536
3241
3242# the bytes of xs[i - k + 1 .. i] onto acc, built from the last number back (no reverse)
3243def chunk_bytes(k: Nat, +i: U32, rx: Array<F32> & F32, acc: List<&2, U32>) -> Array<F32> & List<&2, U32>:
3244 match k rx:
3245 case 0n Tuple{a, v}: Tuple{a, acc}
3246 case 1n+q Tuple{a, v}:
3247 +u = bits(v)
3248 chunk_bytes(q, U32.sub(i, 1), Array.get(F32, a, U32.sub(i, 1)), byte_of(u, 0n) <> byte_of(u, 8n) <> byte_of(u, 16n) <> byte_of(u, 24n) <> acc)
3249
3250def wc_ok(-R: Type, rr: File & Result<&1, &1, U32 & String, Unit>, a: Array<F32>, k: File -> Array<F32> -> IO(R)) -> IO(R):
3251 match rr:
3252 case Tuple{f, res}:
3253 do IO<R>:
3254 u : Unit <- IO.pass(Unit, res)
3255 k(f, a)
3256
3257def wc_bytes(-R: Type, f: File, r: Array<F32> & List<&2, U32>, k: File -> Array<F32> -> IO(R)) -> IO(R):
3258 match r:
3259 case Tuple{a, bs}:
3260 do IO<R>:
3261 rr : File & Result<&1, &1, U32 & String, Unit> <- File.write_bytes(f, bs)
3262 wc_ok(R, rr, a, k)
3263
3264# k chunks of a's numbers lo .. n - 1
3265def write_chunks(k: Nat, f: File, a: Array<F32>, +lo: U32, +n: U32) -> IO(File & Array<F32>):
3266 match k:
3267 case 0n:
3268 IO.pure(File & Array<F32>, Tuple{f, a})
3269 case 1n+q:
3270 +hi = U32.min((lo + WCHUNK() : U32), n)
3271 wc_bytes(File & Array<F32>, f, chunk_bytes(U32.to_nat(U32.sub(hi, lo)), U32.sub(hi, 1), Array.get(F32, a, U32.sub(hi, 1)), []), f2 => a2 => write_chunks(q, f2, a2, hi, n))
3272
3273def wf_close(r: File & Array<F32>) -> IO(Array<F32>):
3274 match r:
3275 case Tuple{f, a}:
3276 do IO<Array<F32>>:
3277 File.close(f)
3278 return a
3279
3280# cells 0 .. n - 1 of a, as little-endian F32, into a new file at path; a comes back
3281def write_f32(+path: String, a: Array<F32>, +n: U32) -> IO(Array<F32>):
3282 do IO<Array<F32>>:
3283 f : File <- IO.try(File, File.open(path, "w"))
3284 r : File & Array<F32> <- write_chunks(U32.to_nat(U32.div((n + U32.sub(WCHUNK(), 1) : U32), WCHUNK())), f, a, 0, n)
3285 wf_close(r)
3286
3287def ps_fin(-P: Type, +path: String, +n: U32, r: PA<P>) -> IO(P):
3288 match r:
3289 case PA{p, xs, off}:
3290 do IO<P>:
3291 a : Array<F32> <- write_f32(path, xs, n)
3292 return p
3293
3294# every parameter of p, in dump order, into a file at path (Layer.count numbers), through one array
3295def Params.save(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, p: P) -> IO(P):
3296 +n = Layer.count(P, X, Y, S, m)
3297 ps_fin(P, path, n, Layer.store(P, X, Y, S, m, p, Array.new(F32, depth(U32.to_nat(n)), 0.0), 0))
3298
3299# ---- a training state into a file: the parameters, then the moments in use, then Adam's b1^t and b2^t ----
3300
3301def tsa_fin(-P: Type, +path: String, g: P, st: OptState<P>, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3302 match r:
3303 case PA{p, xs, +off}:
3304 do IO<Train<P>>:
3305 a : Array<F32> <- write_f32(path, xs, off)
3306 return Train{p, g, st, p1, p2, l}
3307
3308def tsa_m(-P: Type, +path: String, p: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3309 match r:
3310 case PA{mm, xs, off}: tsa_fin(P, path, g, OneMom{mm}, p1, p2, l, PA{p, xs, off})
3311
3312def tsa_pw(-P: Type, +path: String, p: P, mm: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3313 match r:
3314 case PA{vv, xs, +off}:
3315 tsa_fin(P, path, g, TwoMom{mm, vv}, p1, p2, l, PA{p, Array.set(F32, Array.set(F32, xs, off, p1), (off + 1 : U32), p2), (off + 2 : U32)})
3316
3317def tsa_v(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, p: P, vv: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3318 match r:
3319 case PA{mm, xs, off}: tsa_pw(P, path, p, mm, g, p1, p2, l, Layer.store(P, X, Y, S, m, vv, xs, off))
3320
3321def tsa_mm(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, mm: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3322 match r:
3323 case PA{p, xs, off}: tsa_m(P, path, p, g, p1, p2, l, Layer.store(P, X, Y, S, m, mm, xs, off))
3324
3325def tsa_mv(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, mm: P, vv: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3326 match r:
3327 case PA{p, xs, off}: tsa_v(~P, ~X, ~Y, ~S, ~m, path, p, vv, g, p1, p2, l, Layer.store(P, X, Y, S, m, mm, xs, off))
3328
3329def tsp3(-P: Type, +path: String, p: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PTA) -> IO(Train<P>):
3330 match r:
3331 case PTA{t, xs, +off}:
3332 tsa_fin(P, path, g, Packed{t}, p1, p2, l, PA{p, Array.set(F32, Array.set(F32, xs, off, p1), (off + 1 : U32), p2), (off + 2 : U32)})
3333
3334def tsp2(-P: Type, +path: String, g: P, +p1: F32, +p2: F32, +l: F32, t: PTree, r: PA<P>) -> IO(Train<P>):
3335 match r:
3336 case PA{p, xs, +off}: tsp3(P, path, p, g, p1, p2, l, pt_store(t, xs, off))
3337
3338def tsp(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, p: P, g: P, +p1: F32, +p2: F32, +l: F32, rt: PTU) -> IO(Train<P>):
3339 match rt:
3340 case PTU{t, +n}:
3341 tsp2(P, path, g, p1, p2, l, t, Layer.store(P, X, Y, S, m, p, Array.new(F32, depth(U32.to_nat(((Layer.count(P, X, Y, S, m) + n : U32) + 2 : U32))), 0.0), 0))
3342
3343# the layout Train.dump gives, written from one array instead of a list (the same bytes)
3344def Train.save(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, tr: Train<P>) -> IO(Train<P>):
3345 match tr:
3346 case Train{p, g, NoMom{}, +p1, +p2, +l}:
3347 tsa_fin(P, path, g, NoMom{}, p1, p2, l, Layer.store(P, X, Y, S, m, p, Array.new(F32, depth(U32.to_nat(Layer.count(P, X, Y, S, m))), 0.0), 0))
3348 case Train{p, g, OneMom{mm}, +p1, +p2, +l}:
3349 tsa_mm(~P, ~X, ~Y, ~S, ~m, path, mm, g, p1, p2, l, Layer.store(P, X, Y, S, m, p, Array.new(F32, depth(U32.to_nat((2 * Layer.count(P, X, Y, S, m) : U32))), 0.0), 0))
3350 case Train{p, g, TwoMom{mm, vv}, +p1, +p2, +l}:
3351 tsa_mv(~P, ~X, ~Y, ~S, ~m, path, mm, vv, g, p1, p2, l, Layer.store(P, X, Y, S, m, p, Array.new(F32, depth(U32.to_nat(((3 * Layer.count(P, X, Y, S, m) : U32) + 2 : U32))), 0.0), 0))
3352 case Train{p, g, Packed{t}, +p1, +p2, +l}:
3353 tsp(~P, ~X, ~Y, ~S, ~m, path, p, g, p1, p2, l, pt_cells(t))