main.bend

bendtorch@1.2.0.0 · 3,423 lines · 601 definitions

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