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