main.bend

bendtorch@1.3.0.0 · 3,707 lines · 645 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.18.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  # bend-lint: disable-next bend-lint-rules/performance/gpu_fit -- zeros are made once per run (Train.new, OptState.new), and pair_zeros already forks not-expected-because: the combinators nest, so a bang here would run inside a bang
845  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)}
846
847# Owned layer composition for dimensions read at run time. Each child record is opened once; each closure goes to
848# exactly one corresponding closure of the composite layer. The template-based Seq above keeps its existing laws.
849def 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>>:
850  match r:
851    case Fwd{p1, y, s1}: seq_f3(P1, P2, Z, S1, S2, p1, s1, fb(p2, y))
852
853def 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>>:
854  match p:
855    case Two{p1, p2}: seq_at_f2(P1, P2, Y, Z, S1, S2, p2, fb, fa(p1, x))
856
857def 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>:
858  match r:
859    case Bwd{p2, dy, g2}: seq_b3(P1, P2, X, p2, g2, ba(p1, s1, dy, g1))
860
861def 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>:
862  match p s g:
863    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))
864
865def 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>>:
866  match p g st:
867    case Two{p1, p2} Two{g1, g2} Two{s1, s2}: pair_s3(P1, P2, sa(op, p1, g1, s1), sb(op, p2, g2, s2))
868
869def seq_at_dump(-P1: Type, -P2: Type, da: P1 -> PL<P1>, db: P2 -> PL<P2>, p: Two<P1, P2>) -> PL<Two<P1, P2>>:
870  match p:
871    case Two{p1, p2}: pl_pair(P1, P2, da(p1), db(p2))
872
873def seq_at_load2(-P1: Type, -P2: Type, lb: P2 -> Array<F32> -> U32 -> PA<P2>, p2: P2, r: PA<P1>) -> PA<Two<P1, P2>>:
874  match r:
875    case PA{p1, xs, off}: pa_second(P1, P2, p1, lb(p2, xs, off))
876
877def 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>>:
878  match p:
879    case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off))
880
881def 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>>:
882  match p:
883    case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off))
884
885def 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>>:
886  match a b:
887    case Layer{fa, ba, sa, za, da, la, ta, na} Layer{fb, bb, sb, zb, db, lb, tb, nb}:
888      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)}
889
890# =====================================================================================================
891# Losses: the mean loss over the batch, and its gradient with respect to the output
892# =====================================================================================================
893
894type LossOut<-Y: Type> is Type:
895  LossOut{loss: F32, dy: Y}
896
897# =====================================================================================================
898# Training: one step of any layer under any loss
899# =====================================================================================================
900
901# Parameters, gradient (zero between steps), optimizer-specific buffers, Adam's 1 - b1^t and 1 - b2^t, last loss.
902type Train<-P: Type> is Type:
903  Train{p: P, g: P, st: OptState<P>, p1: F32, p2: F32, loss: F32}
904
905def OptState.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, cfg: OptCfg) -> OptState<P>:
906  match cfg:
907    case CSGD{lr}: NoMom{}
908    case CMomentum{lr, mu}: OneMom{Layer.zeros(P, X, Y, S, m)}
909    case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: Packed{PNew{v16}}
910    case _: TwoMom{Layer.zeros(P, X, Y, S, m), Layer.zeros(P, X, Y, S, m)}
911
912# 1 - b^t starts at 0; a configuration OptCfg.ok rejects starts at NaN, so every loss after the first is NaN.
913def q_start(ok: Bool) -> F32:
914  match ok:
915    case True{}:
916      0.0
917    case False{}:
918      (0.0 / 0.0 : F32)
919
920def Train.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +cfg: OptCfg, p: P) -> Train<P>:
921  +q = q_start(OptCfg.ok(cfg))
922  Train{p, Layer.zeros(P, X, Y, S, m), OptState.new(~P, ~X, ~Y, ~S, ~m, cfg), q, q, 0.0}
923
924def ts5(-P: Type, +l: F32, +p1: F32, +p2: F32, r: Step<P>) -> Train<P>:
925  match r:
926    case Step{p, g, st}:
927      Train{p, g, st, p1, p2, l}
928
929def 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>:
930  match r:
931    case Bwd{p, dx, g}:
932      +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32)
933      +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32)
934      ts5(P, l, q1, q2, Layer.step(P, X, Y, S, m, Opt.at(cfg, q1, q2), p, g, st))
935
936def 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>:
937  match r:
938    case LossOut{l, dy}:
939      ts4(~P, ~X, ~Y, ~S, ~m, cfg, st, p1, p2, l, Layer.bwd(P, X, Y, S, m, p, s, dy, g))
940
941def 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>:
942  match r:
943    case Fwd{p, y, s}:
944      ts3(~P, ~X, ~Y, ~S, ~m, cfg, p, s, g, st, p1, p2, loss(y, t))
945
946# forward, loss, backward and one optimizer step: the same code for every model
947def 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>:
948  match tr:
949    case Train{p, g, st, p1, p2, l}:
950      ts2(~P, ~X, ~Y, ~S, ~T, ~m, ~loss, cfg, g, st, p1, p2, t, Layer.fwd(P, X, Y, S, m, p, x))
951
952# One owned model, including dimensions captured at run time. Its fwd, bwd and step closures each run once.
953def 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>:
954  match r:
955    case Bwd{p, dx, g}:
956      +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32)
957      +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32)
958      ts5(P, l, q1, q2, sf(Opt.at(cfg, q1, q2), p, g, st))
959
960def 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>:
961  match r:
962    case LossOut{l, dy}: train_at_b(P, X, sf, cfg, st, p1, p2, l, bf(p, s, dy, g))
963
964def 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>:
965  match r:
966    case Fwd{p, y, s}: train_at_loss(P, X, Y, S, bf, sf, cfg, p, s, g, st, p1, p2, loss(y, t))
967
968def 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>:
969  match m tr:
970    case Layer{ff, bf, sf, zeros, dump, load, store, count} Train{p, g, st, p1, p2, l}:
971      train_at_f(P, X, Y, S, Target, bf, sf, loss, cfg, g, st, p1, p2, t, ff(p, x))
972
973# =====================================================================================================
974# Reading and printing numbers (bend-ml's and llama.bend's loaders, MIT)
975# =====================================================================================================
976
977# U32 and F32 hold the same Word(32n): these are bitcasts
978def bitcast(u: U32) -> F32:
979  match u:
980    case U32{w}:
981      F32{w}
982
983def bits(x: F32) -> U32:
984  match x:
985    case F32{w}:
986      U32{w}
987
988def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32:
989  bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32))
990
991def f32s(bs: List<&2, U32>) -> List<&2, F32>:
992  match bs:
993    case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}:
994      f32_le(b0, b1, b2, b3) <> f32s(t)
995    case _:
996      Nil{}
997
998def rd_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>):
999  match m:
1000    case Tuple{f, r}:
1001      do IO<List<&2, U32>>:
1002        bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1003        File.close(f)
1004        return bs
1005
1006# n bytes from offset off
1007def read_at(path: String, +off: U32, +n: U32) -> IO(List<&2, U32>):
1008  do IO<List<&2, U32>>:
1009    f : File <- IO.try(File, File.open(path, "r"))
1010    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, off, n)
1011    rd_done(m)
1012
1013def mat_or_die(-r: Nat, -c: Nat, +path: String, m: Maybe<&1, TA.Mat<r, c>>) -> IO(TA.Mat<r, c>):
1014  match m:
1015    case None{}:
1016      IO.die(TA.Mat<r, c>, 1, path ++ ": wrong number of values for this shape")
1017    case Some{x}:
1018      IO.pure(TA.Mat<r, c>, x)
1019
1020# an r x c matrix stored as r*c little-endian F32 from byte off of a file
1021def Mat.load(+r: Nat, +c: Nat, +path: String, +off: U32) -> IO(TA.Mat<r, c>):
1022  do IO<TA.Mat<r, c>>:
1023    bs : List<&2, U32> <- read_at(path, off, (U32.from_nat(Nat.mul(r, c)) * 4 : U32))
1024    mat_or_die(r, c, path, TA.Mat.of_list(r, c, f32s(bs)))
1025
1026def words(xs: List<&2, F32>) -> String:
1027  match xs:
1028    case Nil{}:
1029      ""
1030    case Con{h, t}:
1031      " " ++ U32.show(bits(h)) ++ words(t)
1032
1033def ml_words(-r: Nat, -c: Nat, m: TA.ML<r, c>) -> String:
1034  match m:
1035    case TA.ML{d, xs}:
1036      words(xs)
1037
1038# one line: the name, then each number as the decimal of its 32 bits (exact, for the tests in reference/)
1039def Mat.print(+r: Nat, +c: Nat, +name: String, m: TA.Mat<r, c>) -> IO(Unit):
1040  IO.print(name ++ ml_words(r, c, TA.Mat.to_list(r, c, m)))
1041
1042def F32.print(+name: String, x: F32) -> IO(Unit):
1043  IO.print(name ++ " " ++ U32.show(bits(x)))
1044
1045# n zeros, in an array with room for them
1046def zeros_a(+n: Nat) -> Array<F32>:
1047  Array.new(F32, depth(n), 0.0)
1048
1049# =====================================================================================================
1050# More element-wise activations
1051# =====================================================================================================
1052
1053# an activation whose derivative needs its input (silu, gelu): saves a copy of the input
1054def 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>>:
1055  match m:
1056    case TA.MM2{x, s}:
1057      Fwd{Unit{}, Ew.map(~f, n, c, x), s}
1058
1059def 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>>:
1060  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}
1061
1062def sigmoid_v(+x: F32) -> F32:
1063  (1.0 / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
1064
1065# PyTorch's sigmoid_backward: dy * (1 - y) * y
1066def sigmoid_d(+dy: F32, +y: F32) -> F32:
1067  ((dy * (1.0 - y : F32) : F32) * y : F32)
1068
1069def Sigmoid(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1070  ActOut(~sigmoid_v, ~sigmoid_d, n, c)
1071
1072def tanh_v(+x: F32) -> F32:
1073  F32.tanh(x)
1074
1075def tanh_d(+dy: F32, +y: F32) -> F32:
1076  (dy * (1.0 - (y * y : F32) : F32) : F32)
1077
1078def Tanh(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1079  ActOut(~tanh_v, ~tanh_d, n, c)
1080
1081# silu(x) = x * sigmoid(x), computed as PyTorch's CPU kernel: x / (1 + exp(-x));
1082# its gradient dy * s * (1 + x * (1 - s)) with s = sigmoid(x)
1083def silu_v(+x: F32) -> F32:
1084  (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
1085
1086def silu_d(+dy: F32, +x: F32) -> F32:
1087  +s = sigmoid_v(x)
1088  ((dy * s : F32) * (1.0 + (x * (1.0 - s : F32) : F32) : F32) : F32)
1089
1090def SiLU(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1091  ActIn(~silu_v, ~silu_d, n, c)
1092
1093# GELU in its tanh form (PyTorch's approximate="tanh"): 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3))).
1094# The exact form, with erf, is GELUErf below.
1095def gelu_inner(+x: F32) -> F32:
1096  (0.7978845608028654 * (x + (0.044715 * ((x * x : F32) * x : F32) : F32) : F32) : F32)
1097
1098def gelu_v(+x: F32) -> F32:
1099  (((0.5 * x : F32) * (1.0 + F32.tanh(gelu_inner(x)) : F32)) : F32)
1100
1101def gelu_d(+dy: F32, +x: F32) -> F32:
1102  +th = F32.tanh(gelu_inner(x))
1103  +left = (0.5 * x : F32)
1104  +dleft = (0.5 * (1.0 + th : F32) : F32)
1105  +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)
1106  (dy * (dleft + dright : F32) : F32)
1107
1108def GELU(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1109  ActIn(~gelu_v, ~gelu_d, n, c)
1110
1111# =====================================================================================================
1112# erf in plain Bend, and GELU in its exact form
1113# =====================================================================================================
1114
1115# 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
1116# additions, multiplications and comparisons only: no libm call, so every target runs the same operations (the emitted C
1117# turns fused multiply-adds off); the C and JavaScript programs give the same bits (CUDA is not checked yet). erf is
1118# odd, erf(-x) = -erf(x), so the work is done on a = |x| and the sign put back, in three kinds of pieces:
1119#   - 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
1120#     x itself, so erf(-0) = -0, and NaN reaches it (every comparison with NaN is false) and comes out NaN;
1121#   - 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)
1122#     rounded to F32 and S a polynomial of degree 7 or 8 whose value is small next to K, so the last addition rounds
1123#     once and carries almost all of the error (a - c is exact: a and c are within a factor of 2);
1124#   - a >= 3.92, infinity included: 1, since erf(a) rounds to 1 in F32 from a = 3.9192 on.
1125# The coefficients are least-squares fits against a 50-digit erf on Chebyshev points, rounded to F32. Measured on all
1126# 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
1127# ulp from it on the rest; NaN for every NaN.
1128
1129# 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
1130def erf_hp(+c: F32, +t: F32, +s: F32) -> F32:
1131  (c + (t * s : F32) : F32)
1132
1133def erf_hm(+c: F32, +t: F32, +s: F32) -> F32:
1134  ((t * s : F32) - c : F32)
1135
1136def erf_small(+x: F32) -> F32:
1137  +z = (x * x : F32)
1138  (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)
1139
1140# [0.84375, 1.25): t = a - 1.0
1141def erf_p0(+t: F32) -> F32:
1142  (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)
1143
1144# [1.25, 1.75): t = a - 1.5
1145def erf_p1(+t: F32) -> F32:
1146  (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)
1147
1148# [1.75, 2.25): t = a - 2.0
1149def erf_p2(+t: F32) -> F32:
1150  (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)
1151
1152# [2.25, 2.75): t = a - 2.5
1153def erf_p3(+t: F32) -> F32:
1154  (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)
1155
1156# [2.75, 3.25): t = a - 3.0
1157def erf_p4(+t: F32) -> F32:
1158  (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)
1159
1160# [3.25, 3.92): t = a - 3.5
1161def erf_p5(+t: F32) -> F32:
1162  (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)
1163
1164def erf_sign(+x: F32, +e: F32) -> F32:
1165  pick(F32.is_lt(x, 0.0), F32.neg(e), e)
1166
1167# the pieces are chosen from the top down, so that only one polynomial is computed (each test is a match on a Bool
1168# parameter); b is "a is at least the piece's lower end"
1169def erf_i0(b: Bool, +x: F32, +a: F32) -> F32:
1170  match b:
1171    case True{}: erf_sign(x, erf_p0((a - 1.0 : F32)))
1172    case False{}: erf_small(x)
1173
1174def erf_i1(b: Bool, +x: F32, +a: F32) -> F32:
1175  match b:
1176    case True{}: erf_sign(x, erf_p1((a - 1.5 : F32)))
1177    case False{}: erf_i0(F32.is_ge(a, 0.84375), x, a)
1178
1179def erf_i2(b: Bool, +x: F32, +a: F32) -> F32:
1180  match b:
1181    case True{}: erf_sign(x, erf_p2((a - 2.0 : F32)))
1182    case False{}: erf_i1(F32.is_ge(a, 1.25), x, a)
1183
1184def erf_i3(b: Bool, +x: F32, +a: F32) -> F32:
1185  match b:
1186    case True{}: erf_sign(x, erf_p3((a - 2.5 : F32)))
1187    case False{}: erf_i2(F32.is_ge(a, 1.75), x, a)
1188
1189def erf_i4(b: Bool, +x: F32, +a: F32) -> F32:
1190  match b:
1191    case True{}: erf_sign(x, erf_p4((a - 3.0 : F32)))
1192    case False{}: erf_i3(F32.is_ge(a, 2.25), x, a)
1193
1194def erf_i5(b: Bool, +x: F32, +a: F32) -> F32:
1195  match b:
1196    case True{}: erf_sign(x, erf_p5((a - 3.5 : F32)))
1197    case False{}: erf_i4(F32.is_ge(a, 2.75), x, a)
1198
1199def erf_top(b: Bool, +x: F32, +a: F32) -> F32:
1200  match b:
1201    case True{}: erf_sign(x, 1.0)
1202    case False{}: erf_i5(F32.is_ge(a, 3.25), x, a)
1203
1204def erf(+x: F32) -> F32:
1205  +a = F32.abs(x)
1206  erf_top(F32.is_ge(a, 3.92), x, a)
1207
1208# GELU in its exact form, PyTorch's default (approximate="none"): gelu(x) = x/2 (1 + erf(x / sqrt 2)), in PyTorch's CPU
1209# 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
1210# pdf = exp(-x^2 / 2) / sqrt(2 pi), computed as PyTorch does: 0.5 (1 + erf(x * 0.70710677)) and
1211# 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
1212# float exp).
1213def gelu_erf_cdf(+x: F32) -> F32:
1214  (0.5 * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32)
1215
1216def gelu_erf_v(+x: F32) -> F32:
1217  ((x * 0.5 : F32) * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32)
1218
1219def gelu_erf_d(+dy: F32, +x: F32) -> F32:
1220  +pdf = (0.3989422804014327 * F32.exp(((x * x : F32) * F32.neg(0.5) : F32)) : F32)
1221  (dy * (gelu_erf_cdf(x) + (x * pdf : F32) : F32) : F32)
1222
1223def GELUErf(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1224  ActIn(~gelu_erf_v, ~gelu_erf_d, n, c)
1225
1226# =====================================================================================================
1227# Dropout: each number is kept with probability 1 - p and scaled by 1 / (1 - p), or set to 0. The mask comes from a
1228# counter-based hash of (seed, step, index), so a run is reproducible and a twin can compute the same mask. The step
1229# counter is the layer's "parameter": the optimizer step only increments it, and checkpoints keep it.
1230# =====================================================================================================
1231
1232type DropP is Data:
1233  DropP{t: U32}
1234
1235# lowbias32 (Chris Wellons): an invertible mix of the 32 bits
1236def mix32(+x: U32) -> U32:
1237  +a = U32.xor(x, U32.shrn(x, 16n))
1238  +b = (a * 2146121005 : U32)
1239  +c = U32.xor(b, U32.shrn(b, 15n))
1240  +d = (c * 2221713035 : U32)
1241  U32.xor(d, U32.shrn(d, 16n))
1242
1243# the factor of number i: 1 / (1 - p) if its uniform u = (h >> 8) / 2^24 is >= p, else 0
1244def drop_k(+key: U32, +i: U32, +p: F32, +scale: F32) -> F32:
1245  pick(F32.is_lt((U32.to_f32(U32.shrn(mix32((key + i : U32)), 8n)) / 16777216.0 : F32), p), 0.0, scale)
1246
1247def drop_l(nleft: Nat, +i: U32, +key: U32, +p: F32, +scale: F32, a: Array<F32>) -> Array<F32>:
1248  match nleft:
1249    case 0n:
1250      a
1251    case 1n+q:
1252      drop_l(q, (i + 1 : U32), key, p, scale, Array.set(F32, a, i, drop_k(key, i, p, scale)))
1253
1254def drop_mask(+n: Nat, +c: Nat, +seed: U32, +t: U32, +p: F32) -> TA.Mat<n, c>:
1255  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)))}
1256
1257def drop_f2(-n: Nat, -c: Nat, +t: U32, r: TA.MPair<n, c>) -> Fwd<DropP, TA.Mat<n, c>, TA.Mat<n, c>>:
1258  match r:
1259    case TA.MPair{y, k}:
1260      Fwd{DropP{t}, y, k}
1261
1262def 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>>:
1263  match d:
1264    case DropP{+t}:
1265      drop_f2(n, c, t, Ew.zip(~f32_mul, n, c, x, drop_mask(n, c, seed, t, p)))
1266
1267def drop_b2(-n: Nat, -c: Nat, d: DropP, g: DropP, r: TA.MPair<n, c>) -> Bwd<DropP, TA.Mat<n, c>>:
1268  match r:
1269    case TA.MPair{dx, k}:
1270      Bwd{d, dx, g}
1271
1272def drop_step(+op: Opt, d: DropP, g: DropP, st: OptState<DropP>) -> Step<DropP>:
1273  match d:
1274    case DropP{t}:
1275      Step{DropP{(t + 1 : U32)}, g, st}
1276
1277def drop_dump(d: DropP) -> PL<DropP>:
1278  match d:
1279    case DropP{+t}:
1280      PL{DropP{t}, [bitcast(t)]}
1281
1282def drop_load_fin(+off: U32, r: Array<F32> & F32) -> PA<DropP>:
1283  match r:
1284    case Tuple{a, v}: PA{DropP{bits(v)}, a, (off + 1 : U32)}
1285
1286def drop_load(d: DropP, xs: Array<F32>, +off: U32) -> PA<DropP>:
1287  drop_load_fin(off, Array.get(F32, xs, off))
1288
1289# the counter, bitcast to one F32 as dump writes it
1290def drop_store(d: DropP, xs: Array<F32>, +off: U32) -> PA<DropP>:
1291  match d:
1292    case DropP{+c}: PA{DropP{c}, Array.set(F32, xs, off, bitcast(c)), (off + 1 : U32)}
1293
1294def Dropout(+n: Nat, +c: Nat, +p: F32, +seed: U32) -> Layer<DropP, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1295  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}
1296
1297# =====================================================================================================
1298# 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
1299# expression keeps its order of operations, the ones the row-wise layers had when they worked on lists of rows, so
1300# they give the same bits (NOTES.md, M6)
1301# =====================================================================================================
1302
1303type D2 is Type:
1304  D2{a: Array<F32>, b: Array<F32>, s: F32}
1305
1306type A2 is Type:
1307  A2{a: Array<F32>, b: Array<F32>}
1308
1309# a[ia]·b[ib] + a[ia + sa]·b[ib + sb] + ... (k terms) added to acc left to right; ra and rb hold the
1310# first pair, already read
1311def sdot(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> D2:
1312  match k ra rb:
1313    case 0n Tuple{a, x} Tuple{b, y}:
1314      D2{a, b, acc}
1315    case 1n+p Tuple{a, +x} Tuple{b, +y}:
1316      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)))
1317
1318def dot0(+k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a: Array<F32>, b: Array<F32>) -> D2:
1319  sdot(k, ia, ib, sa, sb, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib))
1320
1321
1322type AS is Type:
1323  AS{a: Array<F32>, s: F32}
1324
1325# the largest of a[idx] (held in r) and the k numbers after it: max(max(a0, a1), a2) ..., from the first
1326def pmax(k: Nat, +idx: U32, cur: F32, r: Array<F32> & F32) -> AS:
1327  match k r:
1328    case 0n Tuple{a, +x}:
1329      AS{a, F32.max(cur, x)}
1330    case 1n+q Tuple{a, +x}:
1331      pmax(q, (idx + 1 : U32), F32.max(cur, x), Array.get(F32, a, (idx + 1 : U32)))
1332
1333def pmax0(k: Nat, +idx: U32, r: Array<F32> & F32) -> AS:
1334  match k r:
1335    case 0n Tuple{a, +x}:
1336      AS{a, x}
1337    case 1n+q Tuple{a, +x}:
1338      pmax(q, (idx + 1 : U32), x, Array.get(F32, a, (idx + 1 : U32)))
1339
1340# a[j] = exp(a[j] - m) for j = idx..idx + k, their sum added to acc in order
1341def pexp(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array<F32> & F32) -> AS:
1342  match k r:
1343    case 0n Tuple{a, +x}:
1344      +e = F32.exp((x - m : F32))
1345      AS{Array.set(F32, a, idx, e), (acc + e : F32)}
1346    case 1n+q Tuple{a, +x}:
1347      +e = F32.exp((x - m : F32))
1348      pexp(q, (idx + 1 : U32), m, (acc + e : F32), Array.get(F32, Array.set(F32, a, idx, e), (idx + 1 : U32)))
1349
1350# the array and how many numbers the last loop of a softmax divided
1351type PC is Type:
1352  PC{a: Array<F32>, c: Nat}
1353
1354def pc_a(r: PC) -> Array<F32>:
1355  match r:
1356    case PC{a, c}:
1357      a
1358
1359# a[j] = a[j] / s for j = idx..idx + k, counted onto c
1360def pdiv(k: Nat, +idx: U32, +s: F32, c: Nat, r: Array<F32> & F32) -> PC:
1361  match k r:
1362    case 0n Tuple{a, +x}:
1363      PC{Array.set(F32, a, idx, (x / s : F32)), 1n+c}
1364    case 1n+q Tuple{a, +x}:
1365      pdiv(q, (idx + 1 : U32), s, 1n+c, Array.get(F32, Array.set(F32, a, idx, (x / s : F32)), (idx + 1 : U32)))
1366
1367def psm3(k: Nat, +off: U32, r: AS) -> PC:
1368  match r:
1369    case AS{a, +s}:
1370      pdiv(k, off, s, 0n, Array.get(F32, a, off))
1371
1372def psm2(+k: Nat, +off: U32, r: AS) -> PC:
1373  match r:
1374    case AS{a, +m}:
1375      psm3(k, off, pexp(k, off, m, 0.0, Array.get(F32, a, off)))
1376
1377# softmax in place over the n numbers of p from off (n >= 1), as PyTorch: exp(x - max) / sum; with the count of the
1378# numbers it divided (the law causal reads it)
1379def psmc(+n: Nat, +off: U32, p: Array<F32>) -> PC:
1380  psm2(Nat.sub(n, 1n), off, pmax0(Nat.sub(n, 1n), off, Array.get(F32, p, off)))
1381
1382def psm(+n: Nat, +off: U32, p: Array<F32>) -> Array<F32>:
1383  pc_a(psmc(n, off, p))
1384
1385# dst[di..] = src[si..], cnt + 1 numbers; r holds src[si], already read
1386def copy_n(cnt: Nat, +si: U32, +di: U32, o: Array<F32>, r: Array<F32> & F32) -> A2:
1387  match cnt r:
1388    case 0n Tuple{s, +x}:
1389      A2{s, Array.set(F32, o, di, x)}
1390    case 1n+n Tuple{s, +x}:
1391      copy_n(n, (si + 1 : U32), (di + 1 : U32), Array.set(F32, o, di, x), Array.get(F32, s, (si + 1 : U32)))
1392
1393# =====================================================================================================
1394# Row-wise layers: softmax, log-softmax, RMSNorm, LayerNorm (each row of an n x c matrix on its own)
1395# =====================================================================================================
1396
1397# rows i.. of x (c wide, left more after i), each softmaxed in place (psm)
1398def smf_rows(left: Nat, +i: U32, +c: Nat, x: Array<F32>) -> Array<F32>:
1399  match left:
1400    case 0n:
1401      psm(c, (i * U32.from_nat(c) : U32), x)
1402    case 1n+q:
1403      smf_rows(q, (i + 1 : U32), c, psm(c, (i * U32.from_nat(c) : U32), x))
1404
1405# the output, and a copy of it saved for the backward pass
1406def out_saved(-n: Nat, -c: Nat, r: TA.MM2<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1407  match r:
1408    case TA.MM2{a, b}:
1409      Fwd{Unit{}, a, b}
1410
1411def sm_fwd(+n: Nat, +c: Nat, x: TA.Mat<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1412  match x:
1413    case TA.Mat{a}:
1414      out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{smf_rows(Nat.sub(n, 1n), 0, c, a)}))
1415
1416# dy[j] = y[j] * (dy[j] - s), j = idx..idx + k; rd and ry hold dy[idx] and y[idx]
1417def smb_fix(k: Nat, +idx: U32, +s: F32, rd: Array<F32> & F32, ry: Array<F32> & F32) -> A2:
1418  match k rd ry:
1419    case 0n Tuple{d, +g} Tuple{y, +v}:
1420      A2{Array.set(F32, d, idx, (v * (g - s : F32) : F32)), y}
1421    case 1n+q Tuple{d, +g} Tuple{y, +v}:
1422      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)))
1423
1424def smb_row(+c: Nat, +off: U32, r: D2) -> A2:
1425  match r:
1426    case D2{d, y, +s}:
1427      smb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off))
1428
1429# softmax backward, as PyTorch's: dx = y * (dy - sum(dy * y)), row after row, in place over dy
1430def smb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2:
1431  match left st:
1432    case 0n A2{d, y}:
1433      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))
1434    case 1n+q A2{d, y}:
1435      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)))
1436
1437def dx_fin(-n: Nat, -c: Nat, r: A2) -> Bwd<Unit, TA.Mat<n, c>>:
1438  match r:
1439    case A2{d, y}:
1440      Bwd{Unit{}, TA.Mat{d}, Unit{}}
1441
1442def sm_bwd(+n: Nat, +c: Nat, s: TA.Mat<n, c>, dy: TA.Mat<n, c>) -> Bwd<Unit, TA.Mat<n, c>>:
1443  match s dy:
1444    case TA.Mat{y} TA.Mat{d}:
1445      dx_fin(n, c, smb_rows(Nat.sub(n, 1n), 0, c, A2{d, y}))
1446
1447# softmax over each row; saves its output
1448def Softmax(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1449  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}
1450
1451# log_softmax(x) = (x - max) - log(sum(exp(x - max))), as PyTorch's CPU kernel
1452
1453# acc + exp(a[j] - m) for j = idx..idx + k, in order
1454def lsum(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array<F32> & F32) -> AS:
1455  match k r:
1456    case 0n Tuple{a, +x}:
1457      AS{a, (acc + F32.exp((x - m : F32)) : F32)}
1458    case 1n+q Tuple{a, +x}:
1459      lsum(q, (idx + 1 : U32), m, (acc + F32.exp((x - m : F32)) : F32), Array.get(F32, a, (idx + 1 : U32)))
1460
1461# a[j] = (a[j] - m) - l for j = idx..idx + k
1462def lw(k: Nat, +idx: U32, +m: F32, +l: F32, r: Array<F32> & F32) -> Array<F32>:
1463  match k r:
1464    case 0n Tuple{a, +x}:
1465      Array.set(F32, a, idx, ((x - m : F32) - l : F32))
1466    case 1n+q Tuple{a, +x}:
1467      lw(q, (idx + 1 : U32), m, l, Array.get(F32, Array.set(F32, a, idx, ((x - m : F32) - l : F32)), (idx + 1 : U32)))
1468
1469def lsm_r3(+c: Nat, +off: U32, +m: F32, r: AS) -> Array<F32>:
1470  match r:
1471    case AS{a, +s}:
1472      lw(Nat.sub(c, 1n), off, m, F32.log(s), Array.get(F32, a, off))
1473
1474def lsm_r2(+c: Nat, +off: U32, r: AS) -> Array<F32>:
1475  match r:
1476    case AS{a, +m}:
1477      lsm_r3(c, off, m, lsum(Nat.sub(c, 1n), off, m, 0.0, Array.get(F32, a, off)))
1478
1479def lsm_row(+c: Nat, +off: U32, x: Array<F32>) -> Array<F32>:
1480  lsm_r2(c, off, pmax0(Nat.sub(c, 1n), off, Array.get(F32, x, off)))
1481
1482def lsmf_rows(left: Nat, +i: U32, +c: Nat, x: Array<F32>) -> Array<F32>:
1483  match left:
1484    case 0n:
1485      lsm_row(c, (i * U32.from_nat(c) : U32), x)
1486    case 1n+q:
1487      lsmf_rows(q, (i + 1 : U32), c, lsm_row(c, (i * U32.from_nat(c) : U32), x))
1488
1489def lsm_fwd(+n: Nat, +c: Nat, x: TA.Mat<n, c>) -> Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
1490  match x:
1491    case TA.Mat{a}:
1492      out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{lsmf_rows(Nat.sub(n, 1n), 0, c, a)}))
1493
1494# acc + a[j] for j = idx..idx + k, in order
1495def asum(k: Nat, +idx: U32, acc: F32, r: Array<F32> & F32) -> AS:
1496  match k r:
1497    case 0n Tuple{a, +x}:
1498      AS{a, (acc + x : F32)}
1499    case 1n+q Tuple{a, +x}:
1500      asum(q, (idx + 1 : U32), (acc + x : F32), Array.get(F32, a, (idx + 1 : U32)))
1501
1502# dy[j] = dy[j] - exp(y[j]) * s, j = idx..idx + k
1503def lsmb_fix(k: Nat, +idx: U32, +s: F32, rd: Array<F32> & F32, ry: Array<F32> & F32) -> A2:
1504  match k rd ry:
1505    case 0n Tuple{d, +g} Tuple{y, +v}:
1506      A2{Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), y}
1507    case 1n+q Tuple{d, +g} Tuple{y, +v}:
1508      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)))
1509
1510def lsmb_row(+c: Nat, +off: U32, y: Array<F32>, r: AS) -> A2:
1511  match r:
1512    case AS{d, +s}:
1513      lsmb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off))
1514
1515# dx = dy - exp(y) * sum(dy), row after row, in place over dy
1516def lsmb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2:
1517  match left st:
1518    case 0n A2{d, y}:
1519      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))))
1520    case 1n+q A2{d, y}:
1521      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)))))
1522
1523def lsm_bwd(+n: Nat, +c: Nat, s: TA.Mat<n, c>, dy: TA.Mat<n, c>) -> Bwd<Unit, TA.Mat<n, c>>:
1524  match s dy:
1525    case TA.Mat{y} TA.Mat{d}:
1526      dx_fin(n, c, lsmb_rows(Nat.sub(n, 1n), 0, c, A2{d, y}))
1527
1528def LogSoftmax(+n: Nat, +c: Nat) -> Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
1529  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}
1530
1531# ---- RMSNorm: y = g * (x / sqrt(mean(x^2) + eps)); its parameter is the gain g (1 x d) ----
1532
1533# one optimizer step on a layer whose parameter is a single matrix
1534def mat_step_fin(-r: Nat, -c: Nat, q: MatStep<r, c>) -> Step<TA.Mat<r, c>>:
1535  match q:
1536    case MatStep{w, g, st}: Step{w, g, st}
1537
1538def 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>>:
1539  mat_step_fin(r, c, Opt.mat(r, c, op, p, g, st))
1540
1541def g_add_fin(-r: Nat, -c: Nat, r2: TA.MPair<r, c>) -> TA.Mat<r, c>:
1542  match r2:
1543    case TA.MPair{a, b}:
1544      a
1545
1546# a + b, dropping b
1547def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
1548  g_add_fin(r, c, Mat.add(r, c, a, b))
1549
1550def dup_fin(-r: Nat, -c: Nat, m: TA.MM2<r, c>) -> Two<TA.Mat<r, c>, TA.Mat<r, c>>:
1551  match m:
1552    case TA.MM2{a, b}:
1553      Two{a, b}
1554
1555# two copies of a matrix
1556def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> Two<TA.Mat<r, c>, TA.Mat<r, c>>:
1557  dup_fin(r, c, TA.Mat.clone(r, c, m))
1558
1559def ml_list(-r: Nat, -c: Nat, ml: TA.ML<r, c>) -> List<&2, F32>:
1560  match ml:
1561    case TA.ML{m, xs}:
1562      xs
1563
1564# ---- 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
1565# left to right from +0.0, each product in the order the list version multiplied, so the bits are the list version's.
1566# ----
1567
1568# acc + a[i]·a[i] + ... over k + 1 numbers from i, in order; r holds a[i], already read
1569def ssq(k: Nat, +i: U32, acc: F32, r: Array<F32> & F32) -> AS:
1570  match k r:
1571    case 0n Tuple{a, +x}:
1572      AS{a, (acc + (x * x : F32) : F32)}
1573    case 1n+p Tuple{a, +x}:
1574      ssq(p, (i + 1 : U32), (acc + (x * x : F32) : F32), Array.get(F32, a, (i + 1 : U32)))
1575
1576# 1 / rms of a row whose sum of squares is s
1577def rms_inv(+s: F32, +df: F32, +eps: F32) -> F32:
1578  (1.0 / F32.sqrt(((s / df : F32) + eps : F32)) : F32)
1579
1580type R3 is Type:
1581  R3{x: Array<F32>, g: Array<F32>, y: Array<F32>}
1582
1583# y[off + j] = g[j] * (x[off + j] * r), j = jj..jj + k; rx and rg hold x[off + jj] and g[jj]
1584def rms_w(k: Nat, +j: U32, +off: U32, +r: F32, y: Array<F32>, rx: Array<F32> & F32, rg: Array<F32> & F32) -> R3:
1585  match k rx rg:
1586    case 0n Tuple{x, +a} Tuple{g, +b}:
1587      R3{x, g, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32))}
1588    case 1n+p Tuple{x, +a} Tuple{g, +b}:
1589      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)))
1590
1591def rmsf_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, g: Array<F32>, y: Array<F32>, r: AS) -> R3:
1592  match r:
1593    case AS{x, +s}:
1594      rms_w(Nat.sub(d, 1n), 0, off, rms_inv(s, df, eps), y, Array.get(F32, x, off), Array.get(F32, g, 0))
1595
1596def rmsf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R3) -> R3:
1597  match s:
1598    case R3{x, g, y}:
1599      rmsf_row2(d, off, df, eps, g, y, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off)))
1600
1601# rows i.. (left + 1 of them, d wide)
1602def rmsf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R3) -> R3:
1603  match left:
1604    case 0n:
1605      rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1606    case 1n+p:
1607      rmsf_rows(p, (i + 1 : U32), d, df, eps, rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1608
1609def rms_f2(-n: Nat, -d: Nat, s: R3) -> Fwd<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1610  match s:
1611    case R3{x, g, y}:
1612      Fwd{TA.Mat{g}, TA.Mat{y}, TA.Mat{x}}
1613
1614def 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>>:
1615  match p x:
1616    case TA.Mat{g} TA.Mat{xa}:
1617      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))}))
1618
1619# with r = 1/rms, xh = x*r, dxh = dy*g: dx = r * (dxh - xh * mean(dxh * xh)); dg += dy * xh, row after row
1620type R5 is Type:
1621  R5{x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, dg: Array<F32>}
1622
1623type R3S is Type:
1624  R3S{x: Array<F32>, dy: Array<F32>, g: Array<F32>, s: F32}
1625
1626# acc + dxh[j]·xh[j] + ..., in order; rx, rd, rg hold x[off + jj], dy[off + jj], g[jj]
1627def 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:
1628  match k rx rd rg:
1629    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1630      R3S{x, dy, g, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32)}
1631    case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1632      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)))
1633
1634def dg_add(+j: U32, +v: F32, r: Array<F32> & F32) -> Array<F32>:
1635  match r:
1636    case Tuple{dg, +old}:
1637      Array.set(F32, dg, j, (old + v : F32))
1638
1639# dx[off + j] = ((dy[j]·g[j]) - (x[j]·r)·m)·r and dg[j] += dy[j]·(x[j]·r), j = jj..jj + k
1640def 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:
1641  match k rx rd rg:
1642    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1643      +xh = (a * r : F32)
1644      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))}
1645    case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1646      +xh = (a * r : F32)
1647      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)))
1648
1649def rmsb_row3(+d: Nat, +off: U32, +df: F32, +r: F32, dx: Array<F32>, dg: Array<F32>, q: R3S) -> R5:
1650  match q:
1651    case R3S{x, dy, g, +s}:
1652      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))
1653
1654def 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:
1655  match q:
1656    case AS{x, +s}:
1657      +r = rms_inv(s, df, eps)
1658      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)))
1659
1660def rmsb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R5) -> R5:
1661  match s:
1662    case R5{x, dy, g, dx, dg}:
1663      rmsb_row2(d, off, df, eps, dy, g, dx, dg, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off)))
1664
1665def rmsb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R5) -> R5:
1666  match left:
1667    case 0n:
1668      rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1669    case 1n+p:
1670      rmsb_rows(p, (i + 1 : U32), d, df, eps, rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1671
1672def gain_bwd_fin(-n: Nat, -d: Nat, -P: Type, p: P, dx: TA.Mat<n, d>, g: P) -> Bwd<P, TA.Mat<n, d>>:
1673  Bwd{p, dx, g}
1674
1675def rms_b2(+d: Nat, -n: Nat, gacc: TA.Mat<1n, d>, s: R5) -> Bwd<TA.Mat<1n, d>, TA.Mat<n, d>>:
1676  match s:
1677    case R5{x, dy, g, dx, dg}:
1678      gain_bwd_fin(n, d, TA.Mat<1n, d>, TA.Mat{g}, TA.Mat{dx}, Mat.sum2(1n, d, gacc, TA.Mat{dg}))
1679
1680def 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>>:
1681  match p x dy:
1682    case TA.Mat{g} TA.Mat{xa} TA.Mat{da}:
1683      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)}))
1684
1685# RMSNorm over rows of d numbers, with a learned gain (Llama's norm); saves its input
1686def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1687  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)}
1688
1689# ---- LayerNorm: y = (x - mean) / sqrt(var + eps) * g + b, var the biased variance ----
1690
1691type LN<-d: Nat> is Type:
1692  LN{g: TA.Mat<1n, d>, b: TA.Mat<1n, d>}
1693
1694# acc + (a[j] - mu)·(a[j] - mu) for j = idx..idx + k, in order
1695def csq(k: Nat, +idx: U32, +mu: F32, acc: F32, r: Array<F32> & F32) -> AS:
1696  match k r:
1697    case 0n Tuple{a, +x}:
1698      AS{a, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32)}
1699    case 1n+q Tuple{a, +x}:
1700      csq(q, (idx + 1 : U32), mu, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32), Array.get(F32, a, (idx + 1 : U32)))
1701
1702type ASM is Type:
1703  ASM{a: Array<F32>, mu: F32, r: F32}
1704
1705def ln_mr2(+df: F32, +eps: F32, +mu: F32, q: AS) -> ASM:
1706  match q:
1707    case AS{a, +v}:
1708      ASM{a, mu, (1.0 / F32.sqrt(((v / df : F32) + eps : F32)) : F32)}
1709
1710def ln_mr1(+d: Nat, +off: U32, +df: F32, +eps: F32, q: AS) -> ASM:
1711  match q:
1712    case AS{a, +s}:
1713      ln_mr2(df, eps, (s / df : F32), csq(Nat.sub(d, 1n), off, (s / df : F32), 0.0, Array.get(F32, a, off)))
1714
1715# a row's mean and 1 / sqrt(var + eps), as ln_xh and ln_r computed them
1716def ln_mr(+d: Nat, +off: U32, +df: F32, +eps: F32, x: Array<F32>) -> ASM:
1717  ln_mr1(d, off, df, eps, asum(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off)))
1718
1719type R4 is Type:
1720  R4{x: Array<F32>, g: Array<F32>, b: Array<F32>, y: Array<F32>}
1721
1722# y[off + j] = ((x - mu) * r) * g[j] + b[j], j = jj..jj + k
1723def 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:
1724  match k rx rg rb:
1725    case 0n Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}:
1726      R4{x, g, b, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32))}
1727    case 1n+q Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}:
1728      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)))
1729
1730def lnf_row2(+d: Nat, +off: U32, g: Array<F32>, b: Array<F32>, y: Array<F32>, m: ASM) -> R4:
1731  match m:
1732    case ASM{x, +mu, +r}:
1733      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))
1734
1735def lnf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4) -> R4:
1736  match s:
1737    case R4{x, g, b, y}:
1738      lnf_row2(d, off, g, b, y, ln_mr(d, off, df, eps, x))
1739
1740def lnf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4) -> R4:
1741  match left:
1742    case 0n:
1743      lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1744    case 1n+q:
1745      lnf_rows(q, (i + 1 : U32), d, df, eps, lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1746
1747def ln_f2(-n: Nat, -d: Nat, s: R4) -> Fwd<LN<d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1748  match s:
1749    case R4{x, g, b, y}:
1750      Fwd{LN{TA.Mat{g}, TA.Mat{b}}, TA.Mat{y}, TA.Mat{x}}
1751
1752def 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>>:
1753  match p x:
1754    case LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{xa}:
1755      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))}))
1756
1757# with xh = (x - mu) * r and dxh = dy * g: dx = r * ((dxh - mean(dxh)) - xh * mean(dxh * xh)); dg += dy * xh; db += dy
1758type R3SS is Type:
1759  R3SS{x: Array<F32>, dy: Array<F32>, g: Array<F32>, sa: F32, sb: F32}
1760
1761# sa + dxh[j] and sb + dxh[j]·xh[j], j = jj..jj + k, each in order
1762def 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:
1763  match k rx rd rg:
1764    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1765      R3SS{x, dy, g, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32)}
1766    case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1767      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)))
1768
1769def ln_dw(+j: U32, +v: F32, +w: F32, st: A2) -> A2:
1770  match st:
1771    case A2{dg, db}:
1772      A2{dg_add(j, v, Array.get(F32, dg, j)), dg_add(j, w, Array.get(F32, db, j))}
1773
1774type R4B is Type:
1775  R4B{x: Array<F32>, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, w: A2}
1776
1777# dx[off + j] = (((dxh - ma) - xh * mb) * r), dg[j] += dy * xh, db[j] += dy, j = jj..jj + k (ma, mb the two means)
1778def 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:
1779  match k rx rd rg:
1780    case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1781      +xh = ((a - mu : F32) * r : F32)
1782      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)}
1783    case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}:
1784      +xh = ((a - mu : F32) * r : F32)
1785      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)))
1786
1787def lnb_row3(+d: Nat, +off: U32, +df: F32, +mu: F32, +r: F32, dx: Array<F32>, w: A2, q: R3SS) -> R4B:
1788  match q:
1789    case R3SS{x, dy, g, +sa, +sb}:
1790      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))
1791
1792def lnb_row2(+d: Nat, +off: U32, +df: F32, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, w: A2, m: ASM) -> R4B:
1793  match m:
1794    case ASM{x, +mu, +r}:
1795      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)))
1796
1797def lnb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4B) -> R4B:
1798  match s:
1799    case R4B{x, dy, g, dx, w}:
1800      lnb_row2(d, off, df, dy, g, dx, w, ln_mr(d, off, df, eps, x))
1801
1802def lnb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4B) -> R4B:
1803  match left:
1804    case 0n:
1805      lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)
1806    case 1n+q:
1807      lnb_rows(q, (i + 1 : U32), d, df, eps, lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s))
1808
1809def ln_b2(+d: Nat, -n: Nat, b: TA.Mat<1n, d>, gg: LN<d>, s: R4B) -> Bwd<LN<d>, TA.Mat<n, d>>:
1810  match gg s:
1811    case LN{ga, gb} R4B{x, dy, g, dx, A2{dg, db}}:
1812      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})})
1813
1814def 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>>:
1815  match p x dy:
1816    case LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da}:
1817      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)}}))
1818
1819def ln_pair(-d: Nat, p: LN<d>) -> Two<TA.Mat<1n, d>, TA.Mat<1n, d>>:
1820  match p:
1821    case LN{w, b}: Two{w, b}
1822
1823def ln_unpair(-d: Nat, p: Two<TA.Mat<1n, d>, TA.Mat<1n, d>>) -> LN<d>:
1824  match p:
1825    case Two{w, b}: LN{w, b}
1826
1827def ln_s2(-d: Nat, a: MatStep<1n, d>, c: MatStep<1n, d>) -> Step<LN<d>>:
1828  match a c:
1829    case MatStep{w, gw, sw} MatStep{b, gb, sb}:
1830      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))}
1831
1832def 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>>:
1833  match p g st:
1834    case LN{w, b} LN{gw, gb} Two{sw, sb}:
1835      ln_s2(d, Opt.mat(1n, d, op, w, gw, sw), Opt.mat(1n, d, op, b, gb, sb))
1836
1837def ln_step(+d: Nat, +op: Opt, p: LN<d>, g: LN<d>, st: OptState<LN<d>>) -> Step<LN<d>>:
1838  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))))
1839
1840def ln_dump2(-d: Nat, r: PL<Two<TA.Mat<1n, d>, TA.Mat<1n, d>>>) -> PL<LN<d>>:
1841  match r:
1842    case PL{Two{g, b}, xs}:
1843      PL{LN{g, b}, xs}
1844
1845def ln_dump(+d: Nat, p: LN<d>) -> PL<LN<d>>:
1846  match p:
1847    case LN{g, b}:
1848      ln_dump2(d, pl_pair(TA.Mat<1n, d>, TA.Mat<1n, d>, Mat.dump(1n, d, g), Mat.dump(1n, d, b)))
1849
1850def ln_load3(-d: Nat, r: PA<Two<TA.Mat<1n, d>, TA.Mat<1n, d>>>) -> PA<LN<d>>:
1851  match r:
1852    case PA{Two{g, b}, xs, off}: PA{LN{g, b}, xs, off}
1853
1854def ln_load2(+d: Nat, b: TA.Mat<1n, d>, r: PA<TA.Mat<1n, d>>) -> PA<LN<d>>:
1855  match r:
1856    case PA{g, xs, off}:
1857      ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.restore(1n, d, b, xs, off)))
1858
1859def ln_load(+d: Nat, p: LN<d>, xs: Array<F32>, off: U32) -> PA<LN<d>>:
1860  match p:
1861    case LN{g, b}: ln_load2(d, b, Mat.restore(1n, d, g, xs, off))
1862
1863def ln_store2(+d: Nat, b: TA.Mat<1n, d>, r: PA<TA.Mat<1n, d>>) -> PA<LN<d>>:
1864  match r:
1865    case PA{g, xs, off}:
1866      ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.store(1n, d, b, xs, off)))
1867
1868def ln_store(+d: Nat, p: LN<d>, xs: Array<F32>, off: U32) -> PA<LN<d>>:
1869  match p:
1870    case LN{g, b}: ln_store2(d, b, Mat.store(1n, d, g, xs, off))
1871
1872def LayerNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer<LN<d>, TA.Mat<n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
1873  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))}
1874
1875# ---- GroupNorm: LayerNorm over one image's group of channels, with a gain and a bias per channel ----
1876# Images are channels last: x is (n·hw) x c with c = ng·cg, image s owns rows s·hw .. s·hw + hw - 1 and group k columns
1877# k·cg .. k·cg + cg - 1. Each (image, group) block of hw·cg numbers is normalized by its own mean and biased variance,
1878# then y = xh * g[ch] + b[ch], as nn.GroupNorm. Writing c as ng·cg in the type makes "ng divides c" hold by
1879# construction. A block is walked one row segment at a time with LayerNorm's loops: `off` is the row's first cell and
1880# `j` the group's first channel, so a segment's cells are off + j .. off + j + cg - 1 and its gains g[j .. j + cg - 1].
1881
1882# acc + the segments of the block's remaining `left` rows, each one row (c cells) below the last, in order
1883def gn_sum(left: Nat, +off: U32, +c: U32, +cg: Nat, q: AS) -> AS:
1884  match left q:
1885    case 0n r:
1886      r
1887    case 1n+l AS{a, +s}:
1888      gn_sum(l, (off + c : U32), c, cg, asum(Nat.sub(cg, 1n), (off + c : U32), s, Array.get(F32, a, (off + c : U32))))
1889
1890# acc + (x - mu)·(x - mu) over the same segments
1891def gn_csq(left: Nat, +off: U32, +c: U32, +cg: Nat, +mu: F32, q: AS) -> AS:
1892  match left q:
1893    case 0n r:
1894      r
1895    case 1n+l AS{a, +s}:
1896      gn_csq(l, (off + c : U32), c, cg, mu, csq(Nat.sub(cg, 1n), (off + c : U32), mu, s, Array.get(F32, a, (off + c : U32))))
1897
1898def gn_mr1(+hw: Nat, +off: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, q: AS) -> ASM:
1899  match q:
1900    case AS{a, +s}:
1901      ln_mr2(nf, eps, (s / nf : F32), gn_csq(Nat.sub(hw, 1n), off, c, cg, (s / nf : F32), csq(Nat.sub(cg, 1n), off, (s / nf : F32), 0.0, Array.get(F32, a, off))))
1902
1903# a block's mean and 1 / sqrt(var + eps); off is its first cell, nf = hw·cg
1904def gn_mr(+hw: Nat, +off: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, x: Array<F32>) -> ASM:
1905  gn_mr1(hw, off, c, cg, nf, eps, gn_sum(Nat.sub(hw, 1n), off, c, cg, asum(Nat.sub(cg, 1n), off, 0.0, Array.get(F32, x, off))))
1906
1907# y over the block's rows, one segment each (LayerNorm's lnf_w with the group's gains)
1908def gnf_rows(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, s: R4) -> R4:
1909  match left s:
1910    case 0n R4{x, g, b, y}:
1911      lnf_w(Nat.sub(cg, 1n), j, off, mu, r, y, Array.get(F32, x, (off + j : U32)), Array.get(F32, g, j), Array.get(F32, b, j))
1912    case 1n+l R4{x, g, b, y}:
1913      gnf_rows(l, (off + c : U32), j, c, cg, mu, r, lnf_w(Nat.sub(cg, 1n), j, off, mu, r, y, Array.get(F32, x, (off + j : U32)), Array.get(F32, g, j), Array.get(F32, b, j)))
1914
1915def gnf_grp2(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, g: Array<F32>, b: Array<F32>, y: Array<F32>, m: ASM) -> R4:
1916  match m:
1917    case ASM{x, +mu, +r}:
1918      gnf_rows(Nat.sub(hw, 1n), off, j, c, cg, mu, r, R4{x, g, b, y})
1919
1920# one (image, group) block: off is the image's first row's first cell, j the group's first channel
1921def gnf_grp(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4) -> R4:
1922  match s:
1923    case R4{x, g, b, y}:
1924      gnf_grp2(hw, off, j, c, cg, g, b, y, gn_mr(hw, (off + j : U32), c, cg, nf, eps, x))
1925
1926def gnf_groups(left: Nat, +off: U32, +j: U32, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4) -> R4:
1927  match left:
1928    case 0n:
1929      gnf_grp(hw, off, j, c, cg, nf, eps, s)
1930    case 1n+q:
1931      gnf_groups(q, off, (j + U32.from_nat(cg) : U32), hw, c, cg, nf, eps, gnf_grp(hw, off, j, c, cg, nf, eps, s))
1932
1933def gnf_images(left: Nat, +off: U32, +ng: Nat, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4) -> R4:
1934  match left:
1935    case 0n:
1936      gnf_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s)
1937    case 1n+q:
1938      gnf_images(q, (off + (U32.from_nat(hw) * c : U32) : U32), ng, hw, c, cg, nf, eps, gnf_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s))
1939
1940def gn_fwd(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32, p: LN<Nat.mul(ng, cg)>, x: TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>) -> Fwd<LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
1941  match p x:
1942    case LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{xa}:
1943      ln_f2(Nat.mul(n, hw), Nat.mul(ng, cg), gnf_images(Nat.sub(n, 1n), 0, ng, hw, U32.from_nat(Nat.mul(ng, cg)), cg, F32.from_nat(Nat.mul(hw, cg)), eps, R4{xa, g, b, zeros_a(Nat.mul(Nat.mul(n, hw), Nat.mul(ng, cg)))}))
1944
1945# the backward pass is LayerNorm's with the block in place of the row: the two means (of dxh and of dxh·xh) are taken
1946# over the whole block, then each cell's dx, dg[ch] += dy·xh and db[ch] += dy, in order
1947def gnb_sums(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, q: R3SS) -> R3SS:
1948  match left q:
1949    case 0n t:
1950      t
1951    case 1n+l R3SS{x, dy, g, +sa, +sb}:
1952      gnb_sums(l, (off + c : U32), j, c, cg, mu, r, lnb_sums(Nat.sub(cg, 1n), j, (off + c : U32), mu, r, sa, sb, Array.get(F32, x, ((off + c : U32) + j : U32)), Array.get(F32, dy, ((off + c : U32) + j : U32)), Array.get(F32, g, j)))
1953
1954def gnb_rows(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, +ma: F32, +mb: F32, s: R4B) -> R4B:
1955  match left s:
1956    case 0n R4B{x, dy, g, dx, w}:
1957      lnb_w(Nat.sub(cg, 1n), j, off, mu, r, ma, mb, dx, w, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j))
1958    case 1n+l R4B{x, dy, g, dx, w}:
1959      gnb_rows(l, (off + c : U32), j, c, cg, mu, r, ma, mb, lnb_w(Nat.sub(cg, 1n), j, off, mu, r, ma, mb, dx, w, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j)))
1960
1961def gnb_grp3(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +mu: F32, +r: F32, dx: Array<F32>, w: A2, q: R3SS) -> R4B:
1962  match q:
1963    case R3SS{x, dy, g, +sa, +sb}:
1964      gnb_rows(Nat.sub(hw, 1n), off, j, c, cg, mu, r, (sa / nf : F32), (sb / nf : F32), R4B{x, dy, g, dx, w})
1965
1966def gnb_grp2(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, dy: Array<F32>, g: Array<F32>, dx: Array<F32>, w: A2, m: ASM) -> R4B:
1967  match m:
1968    case ASM{x, +mu, +r}:
1969      gnb_grp3(hw, off, j, c, cg, nf, mu, r, dx, w, gnb_sums(Nat.sub(hw, 1n), off, j, c, cg, mu, r, lnb_sums(Nat.sub(cg, 1n), j, off, mu, r, 0.0, 0.0, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j))))
1970
1971def gnb_grp(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4B) -> R4B:
1972  match s:
1973    case R4B{x, dy, g, dx, w}:
1974      gnb_grp2(hw, off, j, c, cg, nf, dy, g, dx, w, gn_mr(hw, (off + j : U32), c, cg, nf, eps, x))
1975
1976def gnb_groups(left: Nat, +off: U32, +j: U32, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4B) -> R4B:
1977  match left:
1978    case 0n:
1979      gnb_grp(hw, off, j, c, cg, nf, eps, s)
1980    case 1n+q:
1981      gnb_groups(q, off, (j + U32.from_nat(cg) : U32), hw, c, cg, nf, eps, gnb_grp(hw, off, j, c, cg, nf, eps, s))
1982
1983def gnb_images(left: Nat, +off: U32, +ng: Nat, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4B) -> R4B:
1984  match left:
1985    case 0n:
1986      gnb_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s)
1987    case 1n+q:
1988      gnb_images(q, (off + (U32.from_nat(hw) * c : U32) : U32), ng, hw, c, cg, nf, eps, gnb_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s))
1989
1990def gn_bwd(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32, p: LN<Nat.mul(ng, cg)>, x: TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, dy: TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, gg: LN<Nat.mul(ng, cg)>) -> Bwd<LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
1991  match p x dy:
1992    case LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da}:
1993      ln_b2(Nat.mul(ng, cg), Nat.mul(n, hw), b, gg, gnb_images(Nat.sub(n, 1n), 0, ng, hw, U32.from_nat(Nat.mul(ng, cg)), cg, F32.from_nat(Nat.mul(hw, cg)), eps, R4B{xa, da, g, zeros_a(Nat.mul(Nat.mul(n, hw), Nat.mul(ng, cg))), A2{zeros_a(Nat.mul(ng, cg)), zeros_a(Nat.mul(ng, cg))}}))
1994
1995# GroupNorm over n images of hw pixels and ng groups of cg channels, channels last; its parameters are LayerNorm's
1996def GroupNorm(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32) -> Layer<LN<Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>, TA.Mat<Nat.mul(n, hw), Nat.mul(ng, cg)>>:
1997  Layer{p => x => gn_fwd(n, hw, ng, cg, eps, p, x), p => s => dy => g => gn_bwd(n, hw, ng, cg, eps, p, s, dy, g), op => p => g => st => ln_step(Nat.mul(ng, cg), op, p, g, st), u => LN{TA.Mat.zeros(1n, Nat.mul(ng, cg)), TA.Mat.zeros(1n, Nat.mul(ng, cg))}, p => ln_dump(Nat.mul(ng, cg), p), p => xs => off => ln_load(Nat.mul(ng, cg), p, xs, off), p => xs => off => ln_store(Nat.mul(ng, cg), p, xs, off), u => U32.from_nat(Nat.mul(2n, Nat.mul(ng, cg)))}
1998
1999# =====================================================================================================
2000# Embedding: row i of the output is row ids[i] of the table (v x d); the gradient adds each row back
2001# =====================================================================================================
2002
2003# row i of the output (d numbers) copied from row ids[i] of the table, for the first n ids
2004def emb_cp(left: Nat, ids: List<&2, Nat>, +i: U32, +d: Nat, st: A2) -> A2:
2005  match left ids st:
2006    case 1n+q Con{+id, t} A2{e, o}:
2007      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))))
2008    case _ _ st:
2009      st
2010
2011def 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>>:
2012  match r:
2013    case A2{e, o}:
2014      Fwd{TA.Mat{e}, TA.Mat{o}, ids}
2015
2016def 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>>:
2017  match e:
2018    case TA.Mat{ea}:
2019      emb_f2(n, v, d, ids, emb_cp(n, ids, 0, d, A2{ea, zeros_a(Nat.mul(n, d))}))
2020
2021# g[gi + j] = g[gi + j] + dy[di + j], j = 0..k
2022def emb_add(k: Nat, +gi: U32, +di: U32, rg: Array<F32> & F32, rd: Array<F32> & F32) -> A2:
2023  match k rg rd:
2024    case 0n Tuple{g, +a} Tuple{dy, +b}:
2025      A2{Array.set(F32, g, gi, (a + b : F32)), dy}
2026    case 1n+q Tuple{g, +a} Tuple{dy, +b}:
2027      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)))
2028
2029# each dY row added into the gradient's row ids[i], in order, by index
2030def emb_sc(left: Nat, ids: List<&2, Nat>, +i: U32, +d: Nat, st: A2) -> A2:
2031  match left ids st:
2032    case 1n+q Con{+id, t} A2{g, dy}:
2033      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))))
2034    case _ _ st:
2035      st
2036
2037def emb_b2(-v: Nat, -d: Nat, e: TA.Mat<v, d>, r: A2) -> Bwd<TA.Mat<v, d>, List<&2, Nat>>:
2038  match r:
2039    case A2{g, dy}:
2040      Bwd{e, Nil{}, TA.Mat{g}}
2041
2042def 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>>:
2043  match dy g:
2044    case TA.Mat{da} TA.Mat{ga}:
2045      emb_b2(v, d, e, emb_sc(n, ids, 0, d, A2{ga, da}))
2046
2047# a table of v rows of d numbers, looked up by n ids; an id has no gradient, so dX is the empty list
2048def Embedding(+n: Nat, +v: Nat, +d: Nat) -> Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
2049  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))}
2050
2051# =====================================================================================================
2052# Linear without bias: y = x·Wᵀ (Llama's projections)
2053# =====================================================================================================
2054
2055def 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>>:
2056  match r:
2057    case TA.MMulNT{x, w, y}:
2058      Fwd{w, y, x}
2059
2060def 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>>:
2061  match r:
2062    case TA.MMulTN{dy, x, g}:
2063      Bwd{w, dx, g}
2064
2065def 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>>:
2066  match r:
2067    case TA.MMul{dy, w, dx}:
2068      lnb_b3(n, i, o, w, dx, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, g))
2069
2070def 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>>:
2071  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))}
2072
2073# =====================================================================================================
2074# Combinators: Id, Fork (one input to two layers), Both (two inputs, two layers), Add2, Mul2, Residual
2075# =====================================================================================================
2076
2077def Id(~X: Type) -> Layer<Unit, X, X, Unit>:
2078  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}
2079
2080# one optimizer step on two layers' parameters, of any types
2081def 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>>:
2082  match p g st:
2083    case Two{p1, p2} Two{g1, g2} Two{s1, s2}:
2084      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))
2085
2086def 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>>:
2087  pair_step2(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, op, p, g, OptState.split(P1, P2, st))
2088
2089def 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>>:
2090  match r:
2091    case Fwd{p2, y2, s2}:
2092      Fwd{Two{p1, p2}, Two{y1, y2}, Two{s1, s2}}
2093
2094# runs lb on x2 once la's forward pass is done
2095def 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>>:
2096  match r:
2097    case Fwd{p1, y1, s1}:
2098      pair_f4(P1, P2, Y1, Y2, S1, S2, p1, y1, s1, Layer.fwd(P2, X2, Y2, S2, lb, p2, x2))
2099
2100# ---- Fork: x goes to both layers (copied by `dup`); the two input gradients are summed by `add` ----
2101
2102def 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>>:
2103  match d:
2104    case Two{xa, xb}:
2105      pair_f3(~P2, ~X, ~Y2, ~S2, ~lb, P1, Y1, S1, p2, xb, Layer.fwd(P1, X, Y1, S1, la, p1, xa))
2106
2107def 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>:
2108  match r:
2109    case Bwd{p1, dx1, g1}:
2110      Bwd{Two{p1, p2}, add(dx1, dx2), Two{g1, g2}}
2111
2112def 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>:
2113  match r:
2114    case Bwd{p2, dx2, g2}:
2115      fork_b3(~X, ~add, P1, P2, p2, g2, dx2, Layer.bwd(P1, X, Y1, S1, la, p1, s1, d1, g1))
2116
2117def 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>:
2118  match p s dy g:
2119    case Two{p1, p2} Two{s1, s2} Two{d1, d2} Two{g1, g2}:
2120      fork_b2(~P1, ~X, ~Y1, ~S1, ~add, ~la, P2, p1, s1, d1, g1, Layer.bwd(P2, X, Y2, S2, lb, p2, s2, d2, g2))
2121
2122def 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>>:
2123  match p:
2124    case Two{p1, p2}:
2125      fork_f2(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p1, p2, dup(x))
2126
2127def 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>>:
2128  # bend-lint: disable-next bend-lint-rules/performance/gpu_fit -- zeros are made once per run (Train.new, OptState.new), and pair_zeros already forks not-expected-because: the combinators nest, so a bang here would run inside a bang
2129  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)}
2130
2131# ---- Both: la on the first of a pair, lb on the second ----
2132
2133def 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>>:
2134  match r:
2135    case Bwd{p1, dx1, g1}:
2136      Bwd{Two{p1, p2}, Two{dx1, dx2}, Two{g1, g2}}
2137
2138def 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>>:
2139  match r:
2140    case Bwd{p2, dx2, g2}:
2141      both_b3(P1, P2, X1, X2, p2, g2, dx2, Layer.bwd(P1, X1, Y1, S1, la, p1, s1, d1, g1))
2142
2143def 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>>:
2144  match p s dy g:
2145    case Two{p1, p2} Two{s1, s2} Two{d1, d2} Two{g1, g2}:
2146      both_b2(~P1, ~X1, ~Y1, ~S1, ~la, P2, X2, p1, s1, d1, g1, Layer.bwd(P2, X2, Y2, S2, lb, p2, s2, d2, g2))
2147
2148def 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>>:
2149  match p x:
2150    case Two{p1, p2} Two{x1, x2}:
2151      pair_f3(~P2, ~X2, ~Y2, ~S2, ~lb, P1, Y1, S1, p2, x2, Layer.fwd(P1, X1, Y1, S1, la, p1, x1))
2152
2153def 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>>:
2154  # bend-lint: disable-next bend-lint-rules/performance/gpu_fit -- zeros are made once per run (Train.new, OptState.new), and pair_zeros already forks not-expected-because: the combinators nest, so a bang here would run inside a bang
2155  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)}
2156
2157# ---- Tie: first, mid, last, where first and last share one parameter (tied embeddings: the table is also the output
2158# layer's weight). Forward runs the three in order with the same P for first and last; backward runs last, mid, first,
2159# both adding into the same gradient, as PyTorch does for a tensor used twice. ----
2160
2161type Three<-A: Type, -B: Type, -C: Type> is Type:
2162  Three{a: A, b: B, c: C}
2163
2164def 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>>:
2165  match r:
2166    case Fwd{p, w, s3}:
2167      Fwd{Two{p, q}, w, Three{s1, s2, s3}}
2168
2169def 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>>:
2170  match r:
2171    case Fwd{q, z, s2}:
2172      tie_f4(P, Q, W, S1, S2, S3, q, s1, s2, Layer.fwd(P, Z, W, S3, lc, p, z))
2173
2174def 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>>:
2175  match r:
2176    case Fwd{p, y, s1}:
2177      tie_f3(~P, ~Z, ~W, ~S3, ~lc, Q, S1, S2, p, s1, Layer.fwd(Q, Y, Z, S2, lb, q, y))
2178
2179def 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>>:
2180  match pq:
2181    case Two{p, q}:
2182      tie_f2(~P, ~Q, ~Y, ~Z, ~W, ~S2, ~S3, ~lb, ~lc, S1, q, Layer.fwd(P, X, Y, S1, la, p, x))
2183
2184def tie_b4(-P: Type, -Q: Type, -X: Type, q: Q, gq: Q, r: Bwd<P, X>) -> Bwd<Two<P, Q>, X>:
2185  match r:
2186    case Bwd{p, dx, gp}:
2187      Bwd{Two{p, q}, dx, Two{gp, gq}}
2188
2189def 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>:
2190  match r:
2191    case Bwd{q, dy, gq}:
2192      tie_b4(P, Q, X, q, gq, Layer.bwd(P, X, Y, S1, la, p, s1, dy, gp))
2193
2194def 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>:
2195  match r:
2196    case Bwd{p, dz, gp}:
2197      tie_b3(~P, ~X, ~Y, ~S1, ~la, Q, p, s1, gp, Layer.bwd(Q, Y, Z, S2, lb, q, s2, dz, gq))
2198
2199def 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>:
2200  match pq s g:
2201    case Two{p, q} Three{s1, s2, s3} Two{gp, gq}:
2202      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))
2203
2204# The shared parameter steps once (through first's step; last's must be the same kind of layer for P, e.g. Embedding
2205# and LinearNB over the same v x d matrix), the middle with its own.
2206def 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>>:
2207  # bend-lint: disable-next bend-lint-rules/performance/gpu_fit -- zeros are made once per run (Train.new, OptState.new), and pair_zeros already forks not-expected-because: the combinators nest, so a bang here would run inside a bang
2208  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)}
2209
2210# ---- Add2 and Mul2: element-wise sum and product of a pair of matrices ----
2211
2212def add2_f(+n: Nat, +c: Nat, x: Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> Fwd<Unit, TA.Mat<n, c>, Unit>:
2213  match x:
2214    case Two{a, b}:
2215      Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}}
2216
2217def Add2(+n: Nat, +c: Nat) -> Layer<Unit, Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Unit>:
2218  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}
2219
2220def 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>>>:
2221  match r:
2222    case TA.MPair{y, b}:
2223      Fwd{Unit{}, y, Two{a2, b2}}
2224
2225def 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>>>:
2226  match da db:
2227    case Two{a1, a2} Two{b1, b2}:
2228      mul2_f3(n, c, a2, b2, Ew.zip(~f32_mul, n, c, a1, b1))
2229
2230def 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>>>:
2231  match x:
2232    case Two{a, b}:
2233      mul2_f2(n, c, Mat.dup(n, c, a), Mat.dup(n, c, b))
2234
2235def mpair_a(-n: Nat, -c: Nat, r: TA.MPair<n, c>) -> TA.Mat<n, c>:
2236  match r:
2237    case TA.MPair{a, b}:
2238      a
2239
2240# d(a*b) = (dy*b, dy*a)
2241def 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>>>:
2242  match s d:
2243    case Two{a, b} Two{d1, d2}:
2244      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{}}
2245
2246def 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>>>:
2247  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}
2248
2249# ---- Residual: y = x + f(x), for any layer f from n x c to n x c ----
2250
2251# over any X with a copy (dup), an addition (add) and a layer that adds a pair (join): Seq(Fork(Id, f), join). Residual
2252# is this over matrices; the law residual runs it over the exact model of torch/spec.bend
2253def 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>>:
2254  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)
2255
2256def 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>>:
2257  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)
2258
2259# =====================================================================================================
2260# RoPE (the rotate-half form of Llama and `transformers`): rows are positions (row r is position r mod t), each row
2261# 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
2262# inv[c] = 1 / base^(2c / (2*half)). No parameters, nothing saved: the backward pass rotates by the opposite angle.
2263# =====================================================================================================
2264
2265# 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
2266# operations the list version used, then every pair rotated in place, so the bits are the list version's.
2267
2268type CS is Type:
2269  CS{c: Array<F32>, s: Array<F32>}
2270
2271# the angle of table entry e: inv[c] * position, position e / half, c = e mod half, inv[c] = 1 / base^(2c / (2*half))
2272def rope_ang(+e: U32, +half: U32, +base: F32) -> F32:
2273  ((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)
2274
2275# entries e.. of the tables, left more after e
2276def rope_tab(left: Nat, +e: U32, +half: U32, +base: F32, tb: CS) -> CS:
2277  match left tb:
2278    case 0n CS{cs, sn}:
2279      +ang = rope_ang(e, half, base)
2280      CS{Array.set(F32, cs, e, F32.cos(ang)), Array.set(F32, sn, e, F32.sin(ang))}
2281    case 1n+q CS{cs, sn}:
2282      +ang = rope_ang(e, half, base)
2283      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))})
2284
2285type RS3 is Type:
2286  RS3{x: Array<F32>, c: Array<F32>, s: Array<F32>}
2287
2288# one pair (a = x[i1], b = x[i2]) rotated in place; fwd picks the direction: (a cos - b sin, b cos + a sin) forward,
2289# (a cos + b sin, b cos - a sin) backward, in PyTorch's order of operations for each
2290def rope_p3(fwd: Bool, +i1: U32, +i2: U32, +a: F32, r2: Array<F32> & F32, rc: Array<F32> & F32, rs: Array<F32> & F32) -> RS3:
2291  match fwd r2 rc rs:
2292    case True{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}:
2293      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}
2294    case False{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}:
2295      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}
2296
2297def rope_p2(+fwd: Bool, +i1: U32, +i2: U32, +k: U32, cs: Array<F32>, sn: Array<F32>, r1: Array<F32> & F32) -> RS3:
2298  match r1:
2299    case Tuple{x, +a}:
2300      rope_p3(fwd, i1, i2, a, Array.get(F32, x, i2), Array.get(F32, cs, k), Array.get(F32, sn, k))
2301
2302# 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
2303def rope_one(+e: U32, +fwd: Bool, +hh: U32, +half: U32, +w: U32, +t: U32, st: RS3) -> RS3:
2304  match st:
2305    case RS3{x, cs, sn}:
2306      +r = U32.div(e, hh)
2307      +c = U32.mod(U32.mod(e, hh), half)
2308      +i1 = (((r * w : U32) + (U32.div(U32.mod(e, hh), half) * (half + half : U32) : U32) : U32) + c : U32)
2309      rope_p2(fwd, i1, (i1 + half : U32), ((U32.mod(r, t) * half : U32) + c : U32), cs, sn, Array.get(F32, x, i1))
2310
2311# pairs e.., left more after e
2312def rope_e(left: Nat, +e: U32, +fwd: Bool, +hh: U32, +half: U32, +w: U32, +t: U32, st: RS3) -> RS3:
2313  match left:
2314    case 0n:
2315      rope_one(e, fwd, hh, half, w, t, st)
2316    case 1n+q:
2317      rope_e(q, (e + 1 : U32), fwd, hh, half, w, t, rope_one(e, fwd, hh, half, w, t, st))
2318
2319def rope_fin(-n: Nat, -w: Nat, st: RS3) -> TA.Mat<n, w>:
2320  match st:
2321    case RS3{x, cs, sn}:
2322      TA.Mat{x}
2323
2324def rope_go2(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, x: Array<F32>, tb: CS) -> TA.Mat<n, w>:
2325  match tb:
2326    case CS{cs, sn}:
2327      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}))
2328
2329def 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>:
2330  match m:
2331    case TA.Mat{x}:
2332      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))}))
2333
2334# n rows (b sequences of t positions), w = h heads of 2*half numbers each
2335def 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>:
2336  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}
2337
2338# =====================================================================================================
2339# Causal attention with grouped-query heads: queries n x (h*hd), keys and values n x (kv*hd), h = kv*g; query head
2340# i reads key/value head i / g. Rows are b sequences of t positions; position p sees positions 0..p of its sequence.
2341#   S = (Q·Kᵀ) * scale, P = softmax(S) over the visible positions, O = P·V
2342# The forward pass keeps Q, K, V; the backward pass computes P again from them (recomputation, not storage).
2343# =====================================================================================================
2344
2345# Q, K and V stay in their arrays. Each sequence's t rows are copied out once into arrays of their own, and every
2346# head's probabilities P (t x t) are an array too. Each sum is added in the order the earlier list version added it
2347# (left to right from +0.0), and each row of P is softmaxed in place over the row's visible positions (prow), so the
2348# bits are the list version's.
2349
2350type AttnS is Type:
2351  AttnS{q: Array<F32>, k: Array<F32>, v: Array<F32>}
2352
2353# the key/value head that query head hi reads, with g query heads per key/value head; the forward and the backward pass
2354# both pick it here (the law kv_of_group says which heads share one)
2355def kv_of(+hi: Nat, +g: Nat) -> Nat:
2356  Nat.div(hi, g)
2357
2358# xs reversed onto acc
2359def rcat(xs: List<&2, F32>, acc: List<&2, F32>) -> List<&2, F32>:
2360  match xs:
2361    case Nil{}:
2362      acc
2363    case Con{+x, t}:
2364      rcat(t, x <> acc)
2365
2366type DR is Type:
2367  DR{a: Array<F32>, b: Array<F32>, row: List<&2, F32>}
2368
2369# the dots of one row of a (hd numbers from ia) with consecutive rows of b (hd numbers from ib, each bs further): the
2370# current one r and m more, in order
2371def drow(m: Nat, +hd: Nat, +ia: U32, +ib: U32, +bs: U32, acc: List<&2, F32>, r: D2) -> DR:
2372  match m r:
2373    case 0n D2{a, b, +s}:
2374      DR{a, b, rcat(s <> acc, Nil{})}
2375    case 1n+p D2{a, b, +s}:
2376      drow(p, hd, ia, (ib + bs : U32), bs, s <> acc, dot0(hd, ia, (ib + bs : U32), 1, 1, a, b))
2377
2378def dr_row(r: DR) -> List<&2, F32>:
2379  match r:
2380    case DR{a, b, row}:
2381      row
2382
2383# 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
2384# of k, ks apart, from kc): the one place attention reads keys, i + 1 of them (the law causal_keys)
2385def qrow(+i: Nat, +hd: Nat, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>) -> DR:
2386  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))
2387
2388type QK is Type:
2389  QK{q: Array<F32>, k: Array<F32>, p: Array<F32>}
2390
2391type PW is Type:
2392  PW{p: Array<F32>, n: Nat}
2393
2394# xs scaled by sc, written into p from idx on; n counts the numbers written (from the n given)
2395def pw(xs: List<&2, F32>, +idx: U32, +sc: F32, n: Nat, p: Array<F32>) -> PW:
2396  match xs:
2397    case Nil{}:
2398      PW{p, n}
2399    case Con{x, t}:
2400      pw(t, (idx + 1 : U32), sc, 1n+n, Array.set(F32, p, idx, (x * sc : F32)))
2401
2402type PR is Type:
2403  PR{q: Array<F32>, k: Array<F32>, p: Array<F32>, n: Nat}
2404
2405def pr_n(r: PR) -> Nat:
2406  match r:
2407    case PR{q, k, p, n}:
2408      n
2409
2410def pr_of(q: Array<F32>, k: Array<F32>, r: PC) -> PR:
2411  match r:
2412    case PC{p, c}:
2413      PR{q, k, p, c}
2414
2415# the softmax over the n numbers pw wrote; the row's count is the one the softmax reports
2416def prow3(+off: U32, q: Array<F32>, k: Array<F32>, w: PW) -> PR:
2417  match w:
2418    case PW{p, +n}:
2419      pr_of(q, k, psmc(n, off, p))
2420
2421def prow2(+sc: F32, +off: U32, p: Array<F32>, r: DR) -> PR:
2422  match r:
2423    case DR{q, k, row}:
2424      prow3(off, q, k, pw(row, off, sc, 0n, p))
2425
2426# 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
2427# from the row's start, then the softmax over as many positions as were written: i + 1 (the law causal). The positions
2428# after i are never written and keep their zeros.
2429def 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:
2430  prow2(sc, (U32.from_nat(i) * U32.from_nat(t) : U32), p, qrow(i, hd, qs, qc, ks, kc, q, k))
2431
2432# 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
2433# columns from qc), keys rows of k (ks apart, from kc).
2434def prows(left: Nat, +i: Nat, +t: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, r: PR) -> QK:
2435  match left r:
2436    case 0n PR{q, k, p, n}:
2437      QK{q, k, p}
2438    case 1n+m PR{q, k, p, n}:
2439      prows(m, 1n+i, t, hd, sc, qs, qc, ks, kc, prow(1n+i, t, hd, sc, qs, qc, ks, kc, q, k, p))
2440
2441# P of one head: queries in q (row stride qs, columns from qc), keys in k (row stride ks, columns from kc)
2442def probs(+t: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>) -> QK:
2443  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))))
2444
2445# ---- the same P with every key visible: full attention. `qrow` stops at the
2446# query's own position (i + 1 keys, the causal law); `frow` reads all t keys and
2447# `fprow` softmaxes over all of them, so the row's zeros after position i do not
2448# exist here. Everything else (the scaling, the write order, the softmax) is the
2449# causal path's own functions, so the two agree bit for bit on the triangle the
2450# causal one covers.
2451def frow(+i: Nat, +m: Nat, +hd: Nat, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>) -> DR:
2452  drow(m, 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))
2453
2454def fprow2(+sc: F32, +off: U32, p: Array<F32>, r: DR) -> PR:
2455  prow2(sc, off, p, r)
2456
2457# row i of a head's P with every key: the scaled scores of query i against keys 0..t-1 (frow), then the softmax over
2458# all t of them (the law full)
2459def fprow(+i: Nat, +m: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>, p: Array<F32>) -> PR:
2460  fprow2(sc, (U32.from_nat(i) * U32.from_nat(1n+m) : U32), p, frow(i, m, hd, qs, qc, ks, kc, q, k))
2461
2462def fprows(left: Nat, +i: Nat, +m: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, r: PR) -> QK:
2463  match left r:
2464    case 0n PR{q, k, p, n}:
2465      QK{q, k, p}
2466    case 1n+rest PR{q, k, p, n}:
2467      fprows(rest, 1n+i, m, hd, sc, qs, qc, ks, kc, fprow(1n+i, m, hd, sc, qs, qc, ks, kc, q, k, p))
2468
2469def fprobs(+m: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>) -> QK:
2470  fprows(m, 0n, m, hd, sc, qs, qc, ks, kc, fprow(0n, m, hd, sc, qs, qc, ks, kc, q, k, zeros_a(Nat.mul(1n+m, 1n+m))))
2471
2472type PO is Type:
2473  PO{a: Array<F32>, b: Array<F32>, o: Array<F32>}
2474
2475# 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
2476# on, left more after r (P·V, and dQ = G·K): each sum left to right from +0.0
2477def pmul(left: Nat, +idx: U32, +c: U32, +t: Nat, +bs: U32, +bc: U32, +ow: U32, +oc: U32, o: Array<F32>, r: D2) -> PO:
2478  match left r:
2479    case 0n D2{a, b, +s}:
2480      PO{a, b, Array.set(F32, o, ((U32.div(idx, c) * ow : U32) + (oc + U32.mod(idx, c) : U32) : U32), s)}
2481    case 1n+n D2{a, b, +s}:
2482      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))
2483
2484def acc_set(+idx: U32, +s: F32, r: Array<F32> & F32) -> Array<F32>:
2485  match r:
2486    case Tuple{o, +old}:
2487      Array.set(F32, o, idx, (old + s : F32))
2488
2489# 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
2490# (dK += Gᵀ·Q and dV += Pᵀ·dO): each sum left to right from +0.0, then added into the group's running sum
2491def pacc(left: Nat, +idx: U32, +c: U32, +t: Nat, +bs: U32, +bc: U32, o: Array<F32>, r: D2) -> PO:
2492  match left r:
2493    case 0n D2{a, b, +s}:
2494      PO{a, b, acc_set(idx, s, Array.get(F32, o, idx))}
2495    case 1n+n D2{a, b, +s}:
2496      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))
2497
2498# rows r0.. (t of them, c wide) of src into a new array, src handed back
2499def rows_out(+t: Nat, +c: Nat, +r0: Nat, src: Array<F32>) -> A2:
2500  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))))
2501
2502def a2_b(r: A2) -> Array<F32>:
2503  match r:
2504    case A2{a, b}:
2505      b
2506
2507# a block of t rows, c wide, written into dst from row r0; the block is dropped
2508def rows_in(+t: Nat, +c: Nat, +r0: Nat, blk: Array<F32>, dst: Array<F32>) -> Array<F32>:
2509  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)))
2510
2511def pmul_copy(left: Nat, +idx: U32, +c: U32, +ow: U32, +oc: U32, o: Array<F32>, r: Array<F32> & F32) -> A2:
2512  match left r:
2513    case 0n Tuple{s, +x}:
2514      A2{s, Array.set(F32, o, ((U32.div(idx, c) * ow : U32) + (oc + U32.mod(idx, c) : U32) : U32), x)}
2515    case 1n+n Tuple{s, +x}:
2516      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)))
2517
2518# dst's columns oc.. (c of them, ow wide) of rows 0..t-1 from src (t x c): src dropped
2519def cols_in(+t: Nat, +c: Nat, +ow: U32, +oc: U32, src: Array<F32>, dst: Array<F32>) -> Array<F32>:
2520  a2_b(pmul_copy(Nat.sub(Nat.mul(t, c), 1n), 0, U32.from_nat(c), ow, oc, dst, Array.get(F32, src, 0)))
2521
2522# 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
2523def pick_copy(left: Nat, +idx: U32, +c: U32, +ow: U32, +oc: U32, d: Array<F32>, r: Array<F32> & F32) -> A2:
2524  match left r:
2525    case 0n Tuple{s, +x}:
2526      A2{s, Array.set(F32, d, idx, x)}
2527    case 1n+n Tuple{s, +x}:
2528      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)))
2529
2530# 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
2531def block_out(+t: Nat, +c: Nat, +ow: Nat, +r0: Nat, +c0: Nat, src: Array<F32>) -> A2:
2532  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))))
2533
2534# a t x c block into dst's rows r0.. and columns c0.. (dst ow wide); the block is dropped
2535def block_in(+t: Nat, +c: Nat, +ow: Nat, +r0: Nat, +c0: Nat, blk: Array<F32>, dst: Array<F32>) -> Array<F32>:
2536  cols_in(t, c, U32.from_nat(ow), U32.from_nat(Nat.add(Nat.mul(r0, ow), c0)), blk, dst)
2537
2538# ---- the forward pass of one sequence: every head's O = P·V into its columns of o ----
2539
2540type FS is Type:
2541  FS{q: Array<F32>, k: Array<F32>, v: Array<F32>, o: Array<F32>}
2542
2543type Dm is Data:
2544  Dm{t: Nat, g: Nat, hd: Nat, w: Nat, kvw: Nat, sc: F32}
2545
2546# 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
2547# (kc); both passes take their columns from here (the law kv_cols)
2548type HC is Data:
2549  HC{hc: U32, kc: U32}
2550
2551def head_cols(+dm: Dm, +hi: Nat) -> HC:
2552  match dm:
2553    case Dm{+t, +g, +hd, +w, +kvw, +sc}:
2554      HC{U32.from_nat(Nat.mul(hi, hd)), U32.from_nat(Nat.mul(kv_of(hi, g), hd))}
2555
2556def hc_kc(+c: HC) -> U32:
2557  match c:
2558    case HC{+hc, +kc}:
2559      kc
2560
2561def fs_fin(q: Array<F32>, k: Array<F32>, r: PO) -> FS:
2562  match r:
2563    case PO{p, v, o}:
2564      FS{q, k, v, o}
2565
2566def head_f2(+t: Nat, +hd: Nat, +w: Nat, +kvw: Nat, +hc: U32, +kc: U32, v: Array<F32>, o: Array<F32>, r: QK) -> FS:
2567  match r:
2568    case QK{q, k, p}:
2569      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)))
2570
2571# one head's forward pass at the columns hc (its queries) and kc (its key/value head) that head_cols gives; it has no
2572# head index, so it reads no other columns
2573def head_fwd(+dm: Dm, +cols: HC, s: FS) -> FS:
2574  match dm cols s:
2575    case Dm{+t, +g, +hd, +w, +kvw, +sc} HC{+hc, +kc} FS{q, k, v, o}:
2576      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))
2577
2578def fheads(left: Nat, +hi: Nat, +dm: Dm, s: FS) -> FS:
2579  match left:
2580    case 0n:
2581      s
2582    case 1n+n:
2583      fheads(n, 1n+hi, dm, head_fwd(dm, head_cols(dm, hi), s))
2584
2585def fs_o(s: FS) -> Array<F32>:
2586  match s:
2587    case FS{q, k, v, o}:
2588      o
2589
2590# ---- the backward pass of one sequence: per key/value head, its g query heads' dQ into their columns of dq, their dK
2591# and dV summed in head order (from +0.0), then copied into the head's columns ----
2592
2593type BS is Type:
2594  BS{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, dq: Array<F32>, dka: Array<F32>, dva: Array<F32>}
2595
2596# g[j] = (p[j] * (g[j] - s)) * sc for j = idx..idx + k; rg and rp hold g[idx] and p[idx]
2597def gfix(k: Nat, +idx: U32, +s: F32, +sc: F32, rg: Array<F32> & F32, rp: Array<F32> & F32) -> A2:
2598  match k rg rp:
2599    case 0n Tuple{g, +d} Tuple{p, +y}:
2600      A2{p, Array.set(F32, g, idx, ((y * (d - s : F32) : F32) * sc : F32))}
2601    case 1n+q Tuple{g, +d} Tuple{p, +y}:
2602      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)))
2603
2604def g_set2(+t: Nat, +sc: F32, +off: U32, r: D2) -> A2:
2605  match r:
2606    case D2{g, p, +s}:
2607      gfix(Nat.sub(t, 1n), off, s, sc, Array.get(F32, g, off), Array.get(F32, p, off))
2608
2609type GW is Type:
2610  GW{d: Array<F32>, v: Array<F32>, p: Array<F32>, g: Array<F32>}
2611
2612type GN is Type:
2613  GN{g: Array<F32>, n: Nat}
2614
2615# xs written into g from idx, counted onto n
2616def gw(xs: List<&2, F32>, +idx: U32, n: Nat, g: Array<F32>) -> GN:
2617  match xs:
2618    case Nil{}:
2619      GN{g, n}
2620    case Con{x, t}:
2621      gw(t, (idx + 1 : U32), 1n+n, Array.set(F32, g, idx, x))
2622
2623type GR is Type:
2624  GR{d: Array<F32>, v: Array<F32>, pg: A2, n: Nat}
2625
2626def gr_n(r: GR) -> Nat:
2627  match r:
2628    case GR{d, v, pg, n}:
2629      n
2630
2631# the n positions of the row written: s = sum(dP * P) over them (from +0.0, in order), then G = (P * (dP - s)) * scale
2632def gr3(+sc: F32, +off: U32, d: Array<F32>, v: Array<F32>, p: Array<F32>, w: GN) -> GR:
2633  match w:
2634    case GN{g, +n}:
2635      GR{d, v, g_set2(n, sc, off, dot0(n, off, off, 1, 1, g, p)), n}
2636
2637def gr2(+sc: F32, +off: U32, p: Array<F32>, g: Array<F32>, r: DR) -> GR:
2638  match r:
2639    case DR{d, v, row}:
2640      gr3(sc, off, d, v, p, gw(row, off, 0n, g))
2641
2642# 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
2643# 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
2644# law causal_grad). G starts as zeros and nothing else writes the row, so its positions after i stay 0: no gradient
2645# reaches a key or value after the query's position. (P is 0 there too, so the full row gave the same G.)
2646def 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:
2647  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)))
2648
2649# rows i.. of G; r is row i
2650def grows(left: Nat, +i: Nat, +t: Nat, +hd: Nat, +sc: F32, +ds: U32, +dc: U32, +vs: U32, +vc: U32, r: GR) -> GW:
2651  match left r:
2652    case 0n GR{d, v, A2{p, g}, n}:
2653      GW{d, v, p, g}
2654    case 1n+m GR{d, v, A2{p, g}, n}:
2655      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))
2656
2657def hb5(q: Array<F32>, k: Array<F32>, v: Array<F32>, dq: Array<F32>, dka: Array<F32>, r: PO) -> BS:
2658  match r:
2659    case PO{p, d, dva}:
2660      BS{q, k, v, d, dq, dka, dva}
2661
2662# dK += Gᵀ·Q, then dV += Pᵀ·dO
2663def 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:
2664  match r:
2665    case PO{g, q, dka}:
2666      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)))
2667
2668# dQ = G·K into the head's columns, then dK
2669def 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:
2670  match r:
2671    case PO{g, k, dq}:
2672      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)))
2673
2674def 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:
2675  match r:
2676    case GW{d, v, p, g}:
2677      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)))
2678
2679# P again, then G
2680def 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:
2681  match r:
2682    case QK{q, k, p}:
2683      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)))))
2684
2685# one head's backward pass at head_cols' columns (no head index, as head_fwd)
2686def head_bwd(+dm: Dm, +cols: HC, s: BS) -> BS:
2687  match dm cols s:
2688    case Dm{+t, +g, +hd, +w, +kvw, +sc} HC{+hc, +kc} BS{q, k, v, d, dq, dka, dva}:
2689      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))
2690
2691def bheads(left: Nat, +hi: Nat, +dm: Dm, s: BS) -> BS:
2692  match left:
2693    case 0n:
2694      s
2695    case 1n+n:
2696      bheads(n, 1n+hi, dm, head_bwd(dm, head_cols(dm, hi), s))
2697
2698type BG is Type:
2699  BG{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, dq: Array<F32>, dk: Array<F32>, dv: Array<F32>}
2700
2701# the group's summed dK and dV into key/value head j's columns of dk and dv
2702def bg_fin(+t: Nat, +hd: Nat, +kvw: Nat, +j: Nat, dk: Array<F32>, dv: Array<F32>, s: BS) -> BG:
2703  match s:
2704    case BS{q, k, v, d, dq, dka, dva}:
2705      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)}
2706
2707def bgroup(+dm: Dm, +j: Nat, s: BG) -> BG:
2708  match dm s:
2709    case Dm{+t, +g, +hd, +w, +kvw, +sc} BG{q, k, v, d, dq, dk, dv}:
2710      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))}))
2711
2712def bgroups(left: Nat, +j: Nat, +dm: Dm, s: BG) -> BG:
2713  match left:
2714    case 0n:
2715      s
2716    case 1n+n:
2717      bgroups(n, 1n+j, dm, bgroup(dm, j, s))
2718
2719# ---- the sequences: each one's rows copied out into arrays of its own, the sequences computed in a tree of parallel
2720# lets (each leaf holds its own arrays, nothing shared), the results written back in order ----
2721
2722type Halves<-A: Type> is Type:
2723  Halves{l: List<&1, A>, r: List<&1, A>}
2724
2725# the two halves back as one list, in order (what the laws attn_par_fwd and attn_par_bwd compare the trees with)
2726def hcat(-A: Type, h: Halves<A>) -> List<&1, A>:
2727  match h:
2728    case Halves{l, r}:
2729      List.append(&1, A, l, r)
2730
2731# xs cut after its first n (a tail loop: a helper in the fork tree must not hold the stack)
2732def cut(-A: Type, n: Nat, xs: List<&1, A>, acc: List<&1, A>) -> Halves<A>:
2733  match n xs:
2734    case 1n+m Con{x, xt}:
2735      cut(A, m, xt, x <> acc)
2736    case _ _:
2737      Halves{List.reverse(&1, A, acc), xs}
2738
2739type FL is Type:
2740  FL{q: Array<F32>, k: Array<F32>, v: Array<F32>, xs: List<&1, FS>}
2741
2742# 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
2743# 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
2744# arithmetic (head_cols with j = 0: the law kv_cols), so the units give the bits of the whole sequence computed at once.
2745def fu3(+t: Nat, +gw: Nat, qq: Array<F32>, kk: Array<F32>, xs: List<&1, FS>, q: Array<F32>, k: Array<F32>, r: A2) -> FL:
2746  match r:
2747    case A2{vv, v}:
2748      FL{qq, kk, vv, FS{q, k, v, zeros_a(Nat.mul(t, gw))} <> xs}
2749
2750def 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:
2751  match r:
2752    case A2{kk, k}:
2753      fu3(t, gw, qq, kk, xs, q, k, block_out(t, hd, kvw, r0, kc, vv))
2754
2755def 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:
2756  match r:
2757    case A2{qq, q}:
2758      fu2(t, hd, kvw, r0, kc, gw, qq, vv, xs, q, block_out(t, hd, kvw, r0, kc, kk))
2759
2760def funit(+t: Nat, +w: Nat, +kvw: Nat, +g: Nat, +hd: Nat, +bi: Nat, +j: Nat, s: FL) -> FL:
2761  match s:
2762    case FL{qq, kk, vv, xs}:
2763      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))
2764
2765def fgroups(left: Nat, +j: Nat, +bi: Nat, +t: Nat, +w: Nat, +kvw: Nat, +g: Nat, +hd: Nat, s: FL) -> FL:
2766  match left:
2767    case 0n:
2768      s
2769    case 1n+n:
2770      fgroups(n, 1n+j, bi, t, w, kvw, g, hd, funit(t, w, kvw, g, hd, bi, j, s))
2771
2772# every sequence's units, sequence by sequence and group by group, newest first
2773def fslices(left: Nat, +bi: Nat, +t: Nat, +w: Nat, +kvw: Nat, +kv: Nat, +g: Nat, +hd: Nat, s: FL) -> FL:
2774  match left:
2775    case 0n:
2776      s
2777    case 1n+n:
2778      fslices(n, 1n+bi, t, w, kvw, kv, g, hd, fgroups(kv, 0n, bi, t, w, kvw, g, hd, s))
2779
2780def fleaf(xs: List<&1, FS>, +h: Nat, +dm: Dm, acc: List<&1, Array<F32>>) -> List<&1, Array<F32>>:
2781  match xs:
2782    case Nil{}:
2783      List.reverse(&1, Array<F32>, acc)
2784    case Con{x, rest}:
2785      fleaf(rest, h, dm, fs_o(fheads(h, 0n, dm, x)) <> acc)
2786
2787# the n sequences of gs (cut in two halves) in a tree of parallel lets of depth d; the outputs come back in order
2788def fpar(d: Nat, gs: Halves<FS>, +n: Nat, +h: Nat, +dm: Dm) -> List<&1, Array<F32>>:
2789  match d gs:
2790    case 0n Halves{l, r}:
2791      a b = fleaf(l, h, dm, Nil{}) fleaf(r, h, dm, Nil{})
2792      List.append(&1, Array<F32>, a, b)
2793    case 1n++e Halves{l, r}:
2794      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)
2795      List.append(&1, Array<F32>, a, b)
2796
2797type UO is Type:
2798  UO{xs: List<&1, Array<F32>>, o: Array<F32>}
2799
2800# one sequence's units (t x gw each) into o's rows r0.. at columns 0, gw, 2·gw, ...
2801def ub_row(left: Nat, +j: Nat, +t: Nat, +gw: Nat, +w: Nat, +r0: Nat, r: UO) -> UO:
2802  match left r:
2803    case 1n+n UO{Con{x, rest}, o}:
2804      ub_row(n, 1n+j, t, gw, w, r0, UO{rest, block_in(t, gw, w, r0, Nat.mul(j, gw), x, o)})
2805    case _ r:
2806      r
2807
2808# the units' outputs into o (n x w), in the order fslices made them
2809def blocks_in(left: Nat, +bi: Nat, +t: Nat, +gw: Nat, +w: Nat, +kv: Nat, r: UO) -> Array<F32>:
2810  match left r:
2811    case 0n UO{xs, o}:
2812      o
2813    case 1n+n r:
2814      blocks_in(n, 1n+bi, t, gw, w, kv, ub_row(kv, 0n, t, gw, w, Nat.mul(bi, t), r))
2815
2816type BL is Type:
2817  BL{q: Array<F32>, k: Array<F32>, v: Array<F32>, d: Array<F32>, xs: List<&1, BG>}
2818
2819# the backward pass's unit: one sequence's group j, as fslices' (Q, K, V, and dO's columns of the g heads), with
2820# dQ, dK and dV of its own
2821def 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:
2822  match r:
2823    case A2{dd, d}:
2824      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}
2825
2826def 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:
2827  match r:
2828    case A2{vv, v}:
2829      bu4(t, gw, hd, qq, kk, vv, xs, q, k, v, block_out(t, gw, w, r0, qc, dd))
2830
2831def 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:
2832  match r:
2833    case A2{kk, k}:
2834      bu3(t, w, hd, r0, qc, gw, qq, kk, dd, xs, q, k, block_out(t, hd, kvw, r0, kc, vv))
2835
2836def 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:
2837  match r:
2838    case A2{qq, q}:
2839      bu2(t, w, hd, kvw, r0, qc, kc, gw, qq, vv, dd, xs, q, block_out(t, hd, kvw, r0, kc, kk))
2840
2841def bunit(+t: Nat, +w: Nat, +kvw: Nat, +g: Nat, +hd: Nat, +bi: Nat, +j: Nat, s: BL) -> BL:
2842  match s:
2843    case BL{qq, kk, vv, dd, xs}:
2844      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))
2845
2846def bgroups_u(left: Nat, +j: Nat, +bi: Nat, +t: Nat, +w: Nat, +kvw: Nat, +g: Nat, +hd: Nat, s: BL) -> BL:
2847  match left:
2848    case 0n:
2849      s
2850    case 1n+n:
2851      bgroups_u(n, 1n+j, bi, t, w, kvw, g, hd, bunit(t, w, kvw, g, hd, bi, j, s))
2852
2853def bslices(left: Nat, +bi: Nat, +t: Nat, +w: Nat, +kvw: Nat, +kv: Nat, +g: Nat, +hd: Nat, s: BL) -> BL:
2854  match left:
2855    case 0n:
2856      s
2857    case 1n+n:
2858      bslices(n, 1n+bi, t, w, kvw, kv, g, hd, bgroups_u(kv, 0n, bi, t, w, kvw, g, hd, s))
2859
2860type DB is Type:
2861  DB{dq: Array<F32>, dk: Array<F32>, dv: Array<F32>}
2862
2863def db_of(s: BG) -> DB:
2864  match s:
2865    case BG{q, k, v, d, dq, dk, dv}:
2866      DB{dq, dk, dv}
2867
2868def bleaf(xs: List<&1, BG>, +kv: Nat, +dm: Dm, acc: List<&1, DB>) -> List<&1, DB>:
2869  match xs:
2870    case Nil{}:
2871      List.reverse(&1, DB, acc)
2872    case Con{x, rest}:
2873      bleaf(rest, kv, dm, db_of(bgroups(kv, 0n, dm, x)) <> acc)
2874
2875def bpar(d: Nat, gs: Halves<BG>, +n: Nat, +kv: Nat, +dm: Dm) -> List<&1, DB>:
2876  match d gs:
2877    case 0n Halves{l, r}:
2878      a b = bleaf(l, kv, dm, Nil{}) bleaf(r, kv, dm, Nil{})
2879      List.append(&1, DB, a, b)
2880    case 1n++e Halves{l, r}:
2881      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)
2882      List.append(&1, DB, a, b)
2883
2884type D3 is Type:
2885  D3{dq: Array<F32>, dk: Array<F32>, dv: Array<F32>}
2886
2887type UD is Type:
2888  UD{xs: List<&1, DB>, o: D3}
2889
2890# one sequence's units into rows r0..: dQ at columns j·gw (w wide), dK and dV at j·hd (kvw wide)
2891def ud_row(left: Nat, +j: Nat, +t: Nat, +gw: Nat, +hd: Nat, +w: Nat, +kvw: Nat, +r0: Nat, r: UD) -> UD:
2892  match left r:
2893    case 1n+n UD{Con{DB{bq, bk, bv}, rest}, D3{dq, dk, dv}}:
2894      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)}})
2895    case _ r:
2896      r
2897
2898def dblocks_in(left: Nat, +bi: Nat, +t: Nat, +gw: Nat, +hd: Nat, +w: Nat, +kvw: Nat, +kv: Nat, r: UD) -> D3:
2899  match left r:
2900    case 0n UD{xs, o}:
2901      o
2902    case 1n+n r:
2903      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))
2904
2905# ---- the layer ----
2906
2907type AttnDims is Data:
2908  AttnDims{b: Nat, t: Nat, kv: Nat, g: Nat, hd: Nat, sc: F32}
2909
2910# a unit's dimensions: g query heads (g·hd columns) over one key/value head (hd columns)
2911def dm_of(+dm: AttnDims) -> Dm:
2912  match dm:
2913    case AttnDims{+b, +t, +kv, +g, +hd, +sc}:
2914      Dm{t, g, hd, Nat.mul(g, hd), hd, sc}
2915
2916# 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
2917# big. With the tiny Llama's sequences (t = 16, w = 32) the tree cost more than it gave at 2 sequences (0.48 against
2918# 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,
2919# w = 576, about half a second of a step each) the two sequences one after another kept a single thread busy for a
2920# 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
2921# comparison walks its numbers one unit at a time). The answer never changes the result: the laws attn_par_fwd and
2922# attn_par_bwd show the tree gives the units in order, as one after another.
2923def attn_tree(+b: Nat, +t: Nat, +w: Nat) -> Bool:
2924  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))))
2925
2926def fgo(big: Bool, +b: Nat, +h: Nat, +dm: Dm, xs: List<&1, FS>) -> List<&1, Array<F32>>:
2927  match big:
2928    case True{}:
2929      fpar(Nat.sub(depth(b), 1n), cut(FS, TA.half(b), xs, Nil{}), b, h, dm)
2930    case False{}:
2931      fleaf(xs, h, dm, Nil{})
2932
2933def bgo(big: Bool, +b: Nat, +kv: Nat, +dm: Dm, xs: List<&1, BG>) -> List<&1, DB>:
2934  match big:
2935    case True{}:
2936      bpar(Nat.sub(depth(b), 1n), cut(BG, TA.half(b), xs, Nil{}), b, kv, dm)
2937    case False{}:
2938      bleaf(xs, kv, dm, Nil{})
2939
2940def af2(-n: Nat, +w: Nat, +b: Nat, +kv: Nat, +dm: Dm, s: FL) -> Fwd<Unit, TA.Mat<n, w>, AttnS>:
2941  match dm s:
2942    case Dm{+t, +g, +hd, +gw, +kw, +sc} FL{q, k, v, xs}:
2943      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}}
2944
2945def 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>:
2946  match dm x:
2947    case AttnDims{+b, +t, +kv, +g, +hd, +sc} Two{TA.Mat{q}, Two{TA.Mat{k}, TA.Mat{v}}}:
2948      af2(n, w, b, kv, dm_of(dm), fslices(b, 0n, t, w, kvw, kv, g, hd, FL{q, k, v, Nil{}}))
2949
2950# ---- the same forward with every key visible. The leaf is the only new arithmetic: head_full builds P with fprobs
2951# instead of probs; the tree, the slices and the P·V are the causal path's own functions. Kept beside the causal ones
2952# rather than behind a flag so that no call into Attention changes and the laws attn_par_fwd / attn_par_bwd still
2953# compare the functions they name.
2954def head_full(+dm: Dm, +cols: HC, s: FS) -> FS:
2955  match dm cols s:
2956    case Dm{+t, +g, +hd, +w, +kvw, +sc} HC{+hc, +kc} FS{q, k, v, o}:
2957      head_f2(t, hd, w, kvw, hc, kc, v, o, fprobs(Nat.sub(t, 1n), hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, q, k))
2958
2959def fheads_full(left: Nat, +hi: Nat, +dm: Dm, s: FS) -> FS:
2960  match left:
2961    case 0n:
2962      s
2963    case 1n+n:
2964      fheads_full(n, 1n+hi, dm, head_full(dm, head_cols(dm, hi), s))
2965
2966def fleaf_full(xs: List<&1, FS>, +h: Nat, +dm: Dm, acc: List<&1, Array<F32>>) -> List<&1, Array<F32>>:
2967  match xs:
2968    case Nil{}:
2969      List.reverse(&1, Array<F32>, acc)
2970    case Con{x, rest}:
2971      fleaf_full(rest, h, dm, fs_o(fheads_full(h, 0n, dm, x)) <> acc)
2972
2973def fpar_full(d: Nat, gs: Halves<FS>, +n: Nat, +h: Nat, +dm: Dm) -> List<&1, Array<F32>>:
2974  match d gs:
2975    case 0n Halves{l, r}:
2976      a b = fleaf_full(l, h, dm, Nil{}) fleaf_full(r, h, dm, Nil{})
2977      List.append(&1, Array<F32>, a, b)
2978    case 1n++e Halves{l, r}:
2979      a b = fpar_full(e, cut(FS, TA.half(TA.half(n)), l, Nil{}), TA.half(n), h, dm) fpar_full(e, cut(FS, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), h, dm)
2980      List.append(&1, Array<F32>, a, b)
2981
2982def fgo_full(big: Bool, +b: Nat, +h: Nat, +dm: Dm, xs: List<&1, FS>) -> List<&1, Array<F32>>:
2983  match big:
2984    case True{}:
2985      fpar_full(Nat.sub(depth(b), 1n), cut(FS, TA.half(b), xs, Nil{}), b, h, dm)
2986    case False{}:
2987      fleaf_full(xs, h, dm, Nil{})
2988
2989def af2_full(-n: Nat, +w: Nat, +b: Nat, +kv: Nat, +dm: Dm, s: FL) -> Fwd<Unit, TA.Mat<n, w>, AttnS>:
2990  match dm s:
2991    case Dm{+t, +g, +hd, +gw, +kw, +sc} FL{q, k, v, xs}:
2992      Fwd{Unit{}, TA.Mat{blocks_in(b, 0n, t, gw, w, kv, UO{fgo_full(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}}
2993
2994def attn_fwd_full(+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>:
2995  match dm x:
2996    case AttnDims{+b, +t, +kv, +g, +hd, +sc} Two{TA.Mat{q}, Two{TA.Mat{k}, TA.Mat{v}}}:
2997      af2_full(n, w, b, kv, dm_of(dm), fslices(b, 0n, t, w, kvw, kv, g, hd, FL{q, k, v, Nil{}}))
2998
2999def 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>>>>:
3000  match o:
3001    case D3{dq, dk, dv}:
3002      Bwd{Unit{}, Two{TA.Mat{dq}, Two{TA.Mat{dk}, TA.Mat{dv}}}, Unit{}}
3003
3004def 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>>>>:
3005  match dm s:
3006    case Dm{+t, +g, +hd, +gw, +kw, +sc} BL{q, k, v, d, xs}:
3007      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))}}))
3008
3009def 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>>>>:
3010  match dm s dy:
3011    case AttnDims{+b, +t, +kv, +g, +hd, +sc} AttnS{q, k, v} TA.Mat{d}:
3012      ab2(n, w, kvw, b, kv, dm_of(dm), bslices(b, 0n, t, w, kvw, kv, g, hd, BL{q, k, v, d, Nil{}}))
3013
3014# ---- the backward of the full-visible forward. G's row i also needs every key, so fgrow takes t - 1 dots where grow
3015# takes i; the rest (gfix, pmul, pacc, the group sum, the tree) is the causal path's own code.
3016def fgrow(+i: Nat, +m: 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:
3017  gr2(sc, (U32.from_nat(i) * U32.from_nat(1n+m) : U32), p, g, drow(m, 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)))
3018
3019def fgrows(left: Nat, +i: Nat, +m: Nat, +hd: Nat, +sc: F32, +ds: U32, +dc: U32, +vs: U32, +vc: U32, r: GR) -> GW:
3020  match left r:
3021    case 0n GR{d, v, A2{p, g}, n}:
3022      GW{d, v, p, g}
3023    case 1n+rest GR{d, v, A2{p, g}, n}:
3024      fgrows(rest, 1n+i, m, hd, sc, ds, dc, vs, vc, fgrow(1n+i, m, hd, sc, ds, dc, vs, vc, d, v, p, g))
3025
3026def fhb1(+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:
3027  match r:
3028    case QK{q, k, p}:
3029      hb2(t, hd, w, kvw, hc, kc, q, k, dq, dka, dva, fgrows(Nat.sub(t, 1n), 0n, Nat.sub(t, 1n), hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, fgrow(0n, Nat.sub(t, 1n), hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, d, v, p, zeros_a(Nat.mul(t, t)))))
3030
3031def head_bwd_full(+dm: Dm, +cols: HC, s: BS) -> BS:
3032  match dm cols s:
3033    case Dm{+t, +g, +hd, +w, +kvw, +sc} HC{+hc, +kc} BS{q, k, v, d, dq, dka, dva}:
3034      fhb1(t, hd, w, kvw, sc, hc, kc, v, d, dq, dka, dva, fprobs(Nat.sub(t, 1n), hd, sc, U32.from_nat(w), hc, U32.from_nat(kvw), kc, q, k))
3035
3036def bheads_full(left: Nat, +hi: Nat, +dm: Dm, s: BS) -> BS:
3037  match left:
3038    case 0n:
3039      s
3040    case 1n+n:
3041      bheads_full(n, 1n+hi, dm, head_bwd_full(dm, head_cols(dm, hi), s))
3042
3043def bgroup_full(+dm: Dm, +j: Nat, s: BG) -> BG:
3044  match dm s:
3045    case Dm{+t, +g, +hd, +w, +kvw, +sc} BG{q, k, v, d, dq, dk, dv}:
3046      bg_fin(t, hd, kvw, j, dk, dv, bheads_full(g, Nat.mul(j, g), dm, BS{q, k, v, d, dq, zeros_a(Nat.mul(t, hd)), zeros_a(Nat.mul(t, hd))}))
3047
3048def bgroups_full(left: Nat, +j: Nat, +dm: Dm, s: BG) -> BG:
3049  match left:
3050    case 0n:
3051      s
3052    case 1n+n:
3053      bgroups_full(n, 1n+j, dm, bgroup_full(dm, j, s))
3054
3055def bleaf_full(xs: List<&1, BG>, +kv: Nat, +dm: Dm, acc: List<&1, DB>) -> List<&1, DB>:
3056  match xs:
3057    case Nil{}:
3058      List.reverse(&1, DB, acc)
3059    case Con{x, rest}:
3060      bleaf_full(rest, kv, dm, db_of(bgroups_full(kv, 0n, dm, x)) <> acc)
3061
3062def bpar_full(d: Nat, gs: Halves<BG>, +n: Nat, +kv: Nat, +dm: Dm) -> List<&1, DB>:
3063  match d gs:
3064    case 0n Halves{l, r}:
3065      a b = bleaf_full(l, kv, dm, Nil{}) bleaf_full(r, kv, dm, Nil{})
3066      List.append(&1, DB, a, b)
3067    case 1n++e Halves{l, r}:
3068      a b = bpar_full(e, cut(BG, TA.half(TA.half(n)), l, Nil{}), TA.half(n), kv, dm) bpar_full(e, cut(BG, TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), kv, dm)
3069      List.append(&1, DB, a, b)
3070
3071def bgo_full(big: Bool, +b: Nat, +kv: Nat, +dm: Dm, xs: List<&1, BG>) -> List<&1, DB>:
3072  match big:
3073    case True{}:
3074      bpar_full(Nat.sub(depth(b), 1n), cut(BG, TA.half(b), xs, Nil{}), b, kv, dm)
3075    case False{}:
3076      bleaf_full(xs, kv, dm, Nil{})
3077
3078def ab2_full(+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>>>>:
3079  match dm s:
3080    case Dm{+t, +g, +hd, +gw, +kw, +sc} BL{q, k, v, d, xs}:
3081      ab3(n, w, kvw, dblocks_in(b, 0n, t, gw, hd, w, kvw, kv, UD{bgo_full(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))}}))
3082
3083def attn_bwd_full(+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>>>>:
3084  match dm s dy:
3085    case AttnDims{+b, +t, +kv, +g, +hd, +sc} AttnS{q, k, v} TA.Mat{d}:
3086      ab2_full(n, w, kvw, b, kv, dm_of(dm), bslices(b, 0n, t, w, kvw, kv, g, hd, BL{q, k, v, d, Nil{}}))
3087
3088# the same layer as Attention with every key visible; Attention itself is untouched, so the causal laws and the bits
3089# of every existing model stay as they are
3090def FullAttention(+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>:
3091  Layer{p => x => attn_fwd_full(n, w, kvw, AttnDims{b, t, kv, g, hd, sc}, x), p => s => dy => g2 => attn_bwd_full(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}
3092
3093# the core of attention, after the projections and RoPE; the erased arguments are the shape facts it relies on
3094def 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>:
3095  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}
3096
3097# =====================================================================================================
3098# Losses over arrays (cross-entropy, MSE, NLL)
3099# =====================================================================================================
3100
3101# ---- cross-entropy of the logits z (n x c) against one label per row, mean over the rows (PyTorch's F.cross_entropy,
3102# reduction="mean"); its gradient is (softmax(z) - onehot) / n. Over z's own array, a row at a time, with
3103# bend-ml-tensor-array's Mat.softmax_ce's operations in its order, so its bits. ----
3104
3105def hit(b: Bool) -> F32:
3106  match b:
3107    case True{}:
3108      1.0
3109    case False{}:
3110      0.0
3111
3112# the sum of a[0] .. a[m-1] as a balanced tree of additions (pairwise summation), in place: each pass adds neighbours
3113# (a[i] = a[2i] + a[2i+1], an odd last one carried), until one number is left. Its rounding error grows with log2(m),
3114# not m, as in PyTorch's cascade sum: a running sum of 295k squares was 1e-4 off (NOTES, 2026-10-17)
3115def ts_add(+i: U32, +x: F32, r: Array<F32> & F32) -> Array<F32>:
3116  match r:
3117    case Tuple{a, +y}:
3118      Array.set(F32, a, i, (x + y : F32))
3119
3120def ts_pair(+i: U32, r: Array<F32> & F32) -> Array<F32>:
3121  match r:
3122    case Tuple{a, +x}:
3123      ts_add(i, x, Array.get(F32, a, ((i + i : U32) + 1 : U32)))
3124
3125def ts_pairs(k: Nat, +i: U32, a: Array<F32>) -> Array<F32>:
3126  match k:
3127    case 0n:
3128      a
3129    case 1n+q:
3130      ts_pairs(q, (i + 1 : U32), ts_pair(i, Array.get(F32, a, (i + i : U32))))
3131
3132def ts_put(+i: U32, r: Array<F32> & F32) -> Array<F32>:
3133  match r:
3134    case Tuple{a, +x}:
3135      Array.set(F32, a, i, x)
3136
3137def ts_carry(odd: Bool, +m: U32, a: Array<F32>) -> Array<F32>:
3138  match odd:
3139    case True{}:
3140      ts_put(U32.div(m, 2), Array.get(F32, a, (m - 1 : U32)))
3141    case False{}:
3142      a
3143
3144def ts_pass(+m: U32, a: Array<F32>) -> Array<F32>:
3145  ts_carry(U32.is_eq(U32.mod(m, 2), 1), m, ts_pairs(U32.to_nat(U32.div(m, 2)), 0, a))
3146
3147def ts_first(r: Array<F32> & F32) -> F32:
3148  match r:
3149    case Tuple{a, x}:
3150      x
3151
3152def ts_go(f: Nat, done: Bool, +m: U32, a: Array<F32>) -> F32:
3153  match f done:
3154    case _ True{}:
3155      ts_first(Array.get(F32, a, 0))
3156    case 0n False{}:
3157      ts_first(Array.get(F32, a, 0))
3158    case 1n+g False{}:
3159      ts_go(g, U32.is_le(U32.div((m + 1 : U32), 2), 1), U32.div((m + 1 : U32), 2), ts_pass(m, a))
3160
3161def tree_sum(+m: U32, a: Array<F32>) -> F32:
3162  ts_go(32n, U32.is_le(m, 1), m, a)
3163
3164# Cross-entropy's softmax over a row of c classes: as psm, but the row's exponentials are summed as a balanced tree
3165# (tree_sum, on a copy), since a running F32 sum of SmolLM2's 49152 classes was 3.4 tolerances off PyTorch's loss
3166# (NOTES, M10). psm itself keeps its order: attention's rows are short and the law causal reads its count.
3167# a row's probabilities, its max m and log of its sum of exp(x - m) (log_softmax's pieces)
3168type CSM is Type:
3169  CSM{a: Array<F32>, m: F32, ls: F32}
3170
3171def cesm4(+k: Nat, +off: U32, +m: F32, +s: F32, a: Array<F32>) -> CSM:
3172  CSM{pc_a(pdiv(k, off, s, 0n, Array.get(F32, a, off))), m, F32.log(s)}
3173
3174def cesm3(+k: Nat, +off: U32, +m: F32, r: A2) -> CSM:
3175  match r:
3176    case A2{a, sc}:
3177      cesm4(k, off, m, tree_sum(U32.from_nat(1n+k), sc), a)
3178
3179def cesm_e(+k: Nat, +off: U32, +m: F32, r: AS) -> CSM:
3180  match r:
3181    case AS{a, s}:
3182      cesm3(k, off, m, copy_n(k, off, 0, Array.new(F32, depth(1n+k), 0.0), Array.get(F32, a, off)))
3183
3184def cesm2(+k: Nat, +off: U32, r: AS) -> CSM:
3185  match r:
3186    case AS{a, +m}:
3187      cesm_e(k, off, m, pexp(k, off, m, 0.0, Array.get(F32, a, off)))
3188
3189def ce_sm(+c: Nat, +off: U32, z: Array<F32>) -> CSM:
3190  cesm2(Nat.sub(c, 1n), off, pmax0(Nat.sub(c, 1n), off, Array.get(F32, z, off)))
3191
3192# z[idx] = (p - [j == y]) / nf for j = jj..jj + k, p the row's probabilities
3193def ce_g(k: Nat, +idx: U32, +j: U32, +y: U32, +nf: F32, r: Array<F32> & F32) -> Array<F32>:
3194  match k r:
3195    case 0n Tuple{z, +x}:
3196      Array.set(F32, z, idx, ((x - hit(U32.is_eq(j, y)) : F32) / nf : F32))
3197    case 1n+q Tuple{z, +x}:
3198      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)))
3199
3200# loss - log_softmax(x)[y], log_softmax as PyTorch computes it: (x[y] - m) - log(sum exp(x - m)) (since 2026-10-18; it
3201# was loss - log(max(p[y], 1e-12)), which capped a row's loss at 27.63 where PyTorch's grows, NOTES)
3202def ce_pick2(+xy: F32, loss: F32, r: CSM) -> AS:
3203  match r:
3204    case CSM{z, +m, +ls}:
3205      AS{z, (loss - ((xy - m : F32) - ls : F32) : F32)}
3206
3207def ce_pick3(+c: Nat, +off: U32, loss: F32, r: Array<F32> & F32) -> AS:
3208  match r:
3209    case Tuple{z, +xy}:
3210      ce_pick2(xy, loss, ce_sm(c, off, z))
3211
3212def ce_bad(loss: F32, r: CSM) -> AS:
3213  match r:
3214    case CSM{z, m, ls}:
3215      AS{z, (loss - F32.log(F32.max(0.0, 0.000000000001)) : F32)}
3216
3217# a label past the last class picks probability 0.0, clamped at 1e-12, as tensor-array's pick
3218def ce_pick(ok: Bool, +c: Nat, +off: U32, +idx: U32, loss: F32, z: Array<F32>) -> AS:
3219  match ok:
3220    case True{}:
3221      ce_pick3(c, off, loss, Array.get(F32, z, idx))
3222    case False{}:
3223      ce_bad(loss, ce_sm(c, off, z))
3224
3225def ce_row2(+c: Nat, +off: U32, +y: Nat, +nf: F32, r: AS) -> AS:
3226  match r:
3227    case AS{z, +loss}:
3228      AS{ce_g(Nat.sub(c, 1n), off, 0, U32.from_nat(y), nf, Array.get(F32, z, off)), loss}
3229
3230def ce_row(+c: Nat, +off: U32, +y: Nat, +nf: F32, s: AS) -> AS:
3231  match s:
3232    case AS{z, loss}:
3233      ce_row2(c, off, y, nf, ce_pick(Nat.is_le(1n+y, c), c, off, (off + U32.from_nat(y) : U32), loss, z))
3234
3235def ce_rows(labels: List<&2, Nat>, +c: Nat, +i: U32, +nf: F32, s: AS) -> AS:
3236  match labels:
3237    case Nil{}:
3238      s
3239    case Con{+y, t}:
3240      ce_rows(t, c, (i + 1 : U32), nf, ce_row(c, (i * U32.from_nat(c) : U32), y, nf, s))
3241
3242def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: AS) -> LossOut<TA.Mat<n, c>>:
3243  match r:
3244    case AS{d, +l}:
3245      LossOut{(l / nf : F32), TA.Mat{d}}
3246
3247def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat<n, c>, labels: List<&2, Nat>) -> LossOut<TA.Mat<n, c>>:
3248  match z:
3249    case TA.Mat{d}:
3250      ce_fin(n, c, F32.from_nat(n), ce_rows(labels, c, 0, F32.from_nat(n), AS{d, 0.0}))
3251
3252
3253type D3S is Type:
3254  D3S{y: Array<F32>, t: Array<F32>, q: Array<F32>}
3255
3256# over every number from idx on (k more): d = y - t, q[idx] = d·d, y[idx] = d * kk
3257def mse_l(k: Nat, +idx: U32, +kk: F32, q: Array<F32>, ry: Array<F32> & F32, rt: Array<F32> & F32) -> D3S:
3258  match k ry rt:
3259    case 0n Tuple{y, +a} Tuple{t, +b}:
3260      D3S{Array.set(F32, y, idx, ((a - b : F32) * kk : F32)), t, Array.set(F32, q, idx, ((a - b : F32) * (a - b : F32) : F32))}
3261    case 1n+n Tuple{y, +a} Tuple{t, +b}:
3262      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)))
3263
3264def mse_fin(-n: Nat, -c: Nat, +m: U32, +nc: F32, r: D3S) -> LossOut<TA.Mat<n, c>>:
3265  match r:
3266    case D3S{y, t, q}:
3267      LossOut{(tree_sum(m, q) / nc : F32), TA.Mat{y}}
3268
3269# mean squared error over every number (PyTorch's F.mse_loss); its gradient is 2 (y - t) / (n*c); the squares summed
3270# as a tree
3271def MSE(+n: Nat, +c: Nat, y: TA.Mat<n, c>, t: TA.Mat<n, c>) -> LossOut<TA.Mat<n, c>>:
3272  match y t:
3273    case TA.Mat{ya} TA.Mat{ta}:
3274      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)))
3275
3276type NS is Type:
3277  NS{lp: Array<F32>, g: Array<F32>, s: F32}
3278
3279def nll_hit(+idx: U32, +v: F32, g: Array<F32>, acc: F32, r: Array<F32> & F32) -> NS:
3280  match r:
3281    case Tuple{lp, +x}:
3282      NS{lp, Array.set(F32, g, idx, v), (acc + x : F32)}
3283
3284# a label past the last class picks 0.0 and leaves its gradient row zero, as l_at and onehot do
3285def nll_one(ok: Bool, +idx: U32, +v: F32, acc: F32, lp: Array<F32>, g: Array<F32>) -> NS:
3286  match ok:
3287    case True{}:
3288      nll_hit(idx, v, g, acc, Array.get(F32, lp, idx))
3289    case False{}:
3290      NS{lp, g, (acc + 0.0 : F32)}
3291
3292# the sum of the picked log-probabilities (in row order, from +0.0), and the gradient (-1/n at each label)
3293def nll_rows(left: Nat, labels: List<&2, Nat>, +i: U32, +c: Nat, +v: F32, st: NS) -> NS:
3294  match left labels st:
3295    case 1n+q Con{+y, t} NS{lp, g, acc}:
3296      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))
3297    case _ _ st:
3298      st
3299
3300def nll_fin(-c: Nat, +n: Nat, r: NS) -> LossOut<TA.Mat<n, c>>:
3301  match r:
3302    case NS{lp, g, +s}:
3303      LossOut{F32.neg((s / F32.from_nat(n) : F32)), TA.Mat{g}}
3304
3305# negative log-likelihood of log-probabilities (PyTorch's F.nll_loss, mean over the rows)
3306def NLL(+n: Nat, +c: Nat, logp: TA.Mat<n, c>, labels: List<&2, Nat>) -> LossOut<TA.Mat<n, c>>:
3307  match logp:
3308    case TA.Mat{lp}:
3309      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}))
3310
3311# =====================================================================================================
3312# Checkpoints: parameters, only the buffers in use, and Adam powers, as little-endian F32 (the gradient is zero
3313# between steps and is not saved). A run stopped after step k and resumed from its checkpoint computes the same bits
3314# as a run that never stopped.
3315# =====================================================================================================
3316
3317def byte_of(+u: U32, +k: Nat) -> U32:
3318  U32.and(U32.shrn(u, k), 255)
3319
3320def read_all(+path: String) -> IO(List<&2, U32>):
3321  do IO<List<&2, U32>>:
3322    f : File <- IO.try(File, File.open(path, "r"))
3323    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_bytes(f, 2000000000)
3324    rd_done(m)
3325
3326# ---- packed BF16 moments in a checkpoint (CAdamWBF16) ----
3327# 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:
3328# 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,
3329# 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.
3330
3331def ph_hw(+n: U32) -> U32:
3332  U32.shrn((n + 1 : U32), 1n)
3333
3334def ph_bool(b: Bool) -> U32:
3335  match b:
3336    case True{}: 1
3337    case False{}: 0
3338
3339type PTU is Type:
3340  PTU{t: PTree, n: U32}
3341
3342def ptu_node(ra: PTU, rb: PTU) -> PTU:
3343  match ra rb:
3344    case PTU{a, na} PTU{b, nb}: PTU{PNode{a, b}, ((1 + na : U32) + nb : U32)}
3345
3346# the number of cells a tree takes (the tree comes back)
3347def pt_cells(t: PTree) -> PTU:
3348  match t:
3349    case PNew{v16}: PTU{PNew{v16}, 2}
3350    case PLeaf{+n, m, v}: PTU{PLeaf{n, m, v}, (2 + (2 * ph_hw(n) : U32) : U32)}
3351    case PLeafV{+n, m, v}: PTU{PLeafV{n, m, v}, ((2 + ph_hw(n) : U32) + n : U32)}
3352    # bend-lint: disable-next bend-lint-rules/performance/missed_fork -- a tree of one node per layer, walked once per checkpoint not-expected-because: a fork costs more than the walk
3353    case PNode{a, b}: ptu_node(pt_cells(a), pt_cells(b))
3354
3355type PTA is Type:
3356  PTA{t: PTree, xs: Array<F32>, off: U32}
3357
3358type UA is Type:
3359  UA{a: Array<U32>, xs: Array<F32>}
3360
3361# words i.. of a packed buffer into xs from off, as F32 cells with their bits (both arrays come back)
3362def words_out(left: Nat, +i: U32, +off: U32, ra: Array<U32> & U32, xs: Array<F32>) -> UA:
3363  match left ra:
3364    case 0n Tuple{a, w}: UA{a, xs}
3365    case 1n+q Tuple{a, +w}:
3366      words_out(q, (i + 1 : U32), (off + 1 : U32), Array.get(U32, a, (i + 1 : U32)), Array.set(F32, xs, off, ph_u2f(w)))
3367
3368def words_in(left: Nat, +i: U32, +off: U32, rx: Array<F32> & F32, a: Array<U32>) -> UA:
3369  match left rx:
3370    case 0n Tuple{xs, x}: UA{a, xs}
3371    case 1n+q Tuple{xs, x}:
3372      words_in(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, xs, (off + 1 : U32)), Array.set(U32, a, i, ph_f2u(x)))
3373
3374def pta_then(r: PTA, k: PTree -> Array<F32> -> U32 -> PTA) -> PTA:
3375  match r:
3376    case PTA{t, xs, off}: k(t, xs, off)
3377
3378def pta_node(a: PTree, r: PTA) -> PTA:
3379  match r:
3380    case PTA{b, xs, off}: PTA{PNode{a, b}, xs, off}
3381
3382def pts_v(+n: U32, m: Array<U32>, +off: U32, r: UA) -> PTA:
3383  match r:
3384    case UA{v, xs}: PTA{PLeaf{n, m, v}, xs, (off + ph_hw(n) : U32)}
3385
3386def pts_m(+n: U32, v: Array<U32>, +off: U32, r: UA) -> PTA:
3387  match r:
3388    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))
3389
3390def ptsv_v(+n: U32, m: Array<U32>, +off: U32, r: TA.P2) -> PTA:
3391  match r:
3392    case TA.P2{v, xs}: PTA{PLeafV{n, m, v}, xs, (off + n : U32)}
3393
3394def ptsv_m(+n: U32, v: Array<F32>, +off: U32, r: UA) -> PTA:
3395  match r:
3396    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))
3397
3398# the tree's cells into xs from off
3399def pt_store(t: PTree, xs: Array<F32>, +off: U32) -> PTA:
3400  match t:
3401    case PNew{+v16}:
3402      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)}
3403    case PLeaf{+n, m, v}:
3404      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))))
3405    case PLeafV{+n, m, v}:
3406      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))))
3407    case PNode{a, b}:
3408      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)))
3409
3410# reading a leaf back: fresh buffers filled from xs
3411def ptl_v(+n: U32, m: Array<U32>, +off: U32, r: UA) -> PTA:
3412  match r:
3413    case UA{v, xs}: PTA{PLeaf{n, m, v}, xs, (off + ph_hw(n) : U32)}
3414
3415def ptl_m(+n: U32, +off: U32, r: UA) -> PTA:
3416  match r:
3417    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))))
3418
3419def ptlv_v(+n: U32, m: Array<U32>, +off: U32, r: TA.P2) -> PTA:
3420  match r:
3421    case TA.P2{xs, v}: PTA{PLeafV{n, m, v}, xs, (off + n : U32)}
3422
3423def ptlv_m(+n: U32, +off: U32, r: UA) -> PTA:
3424  match r:
3425    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)))
3426
3427# what a tag says
3428type PK is Data:
3429  KNew{}
3430  KNode{}
3431  KLeaf{}
3432  KLeafV{}
3433
3434def pt_kind2(z: Bool, o: Bool, l: Bool) -> PK:
3435  match z o l:
3436    case True{} _ _: KNew{}
3437    case False{} True{} _: KNode{}
3438    case False{} False{} True{}: KLeaf{}
3439    case False{} False{} False{}: KLeafV{}
3440
3441def pt_kind(+tag: U32) -> PK:
3442  pt_kind2(U32.is_eq(tag, 0), U32.is_eq(tag, 1), U32.is_eq(tag, 2))
3443
3444# the tag at off and the cell after it, to a continuation
3445def pt_next(+off: U32, r: Array<F32> & F32, k: PK -> (Array<F32> & F32) -> PTA) -> PTA:
3446  match r:
3447    case Tuple{xs, x}: k(pt_kind(ph_f2u(x)), Array.get(F32, xs, (off + 1 : U32)))
3448
3449# the part at off, of kind k, with its second cell rn already read; fuel bounds the depth (a model's pairs nest far
3450# less than 1000 deep)
3451def pt_read(fuel: Nat, k: PK, +off: U32, rn: Array<F32> & F32) -> PTA:
3452  match fuel k rn:
3453    case 0n _ Tuple{xs, x}:
3454      PTA{PNew{False{}}, xs, off}
3455    case 1n+f KNew{} Tuple{xs, x}:
3456      PTA{PNew{U32.is_eq(ph_f2u(x), 1)}, xs, (off + 2 : U32)}
3457    case 1n++f KNode{} Tuple{xs, x}:
3458      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))))
3459    case 1n+f KLeaf{} Tuple{xs, x}:
3460      +n = ph_f2u(x)
3461      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))))
3462    case 1n+f KLeafV{} Tuple{xs, x}:
3463      +n = ph_f2u(x)
3464      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))))
3465
3466# a tree from xs at off
3467def pt_load(xs: Array<F32>, +off: U32) -> PTA:
3468  pt_next(off, Array.get(F32, xs, off), k => rn => pt_read(1000n, k, off, rn))
3469
3470type TrainPL<-P: Type> is Type:
3471  TrainPL{tr: Train<P>, xs: List<&2, F32>}
3472
3473def td0(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>) -> TrainPL<P>:
3474  match rp:
3475    case PL{p, xp}: TrainPL{Train{p, g, NoMom{}, p1, p2, l}, xp}
3476
3477def td1(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>, rm: PL<P>) -> TrainPL<P>:
3478  match rp rm:
3479    case PL{p, xp} PL{mm, xm}: TrainPL{Train{p, g, OneMom{mm}, p1, p2, l}, List.append(&2, F32, xp, xm)}
3480
3481def td3(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>, rm: PL<P>, rv: PL<P>) -> TrainPL<P>:
3482  match rp rm rv:
3483    case PL{p, xp} PL{m, xm} PL{v, xv}:
3484      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])))}
3485
3486# the first k cells of an array, in order
3487def cells_l(k: Nat, +i: U32, acc: List<&2, F32>, ra: Array<F32> & F32) -> List<&2, F32>:
3488  match k ra:
3489    case 0n Tuple{a, v}:
3490      List.reverse(&2, F32, acc)
3491    case 1n+q Tuple{a, +v}:
3492      cells_l(q, (i + 1 : U32), v <> acc, Array.get(F32, a, (i + 1 : U32)))
3493
3494def tdp2(-P: Type, p: P, g: P, +p1: F32, +p2: F32, +l: F32, xp: List<&2, F32>, r: PTA) -> TrainPL<P>:
3495  match r:
3496    case PTA{t, xs, +off}:
3497      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]))}
3498
3499def tdp(-P: Type, g: P, +p1: F32, +p2: F32, +l: F32, rp: PL<P>, rt: PTU) -> TrainPL<P>:
3500  match rp rt:
3501    case PL{p, xp} PTU{t, +n}:
3502      tdp2(P, p, g, p1, p2, l, xp, pt_store(t, Array.new(F32, depth(U32.to_nat(n)), 0.0), 0))
3503
3504# every number of a training state: the parameters, then the first moments, the second moments, b1^t and b2^t (packed
3505# moments: their tree's cells, pt_store's layout, in place of the two moments)
3506def Train.dump(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, tr: Train<P>) -> TrainPL<P>:
3507  match tr:
3508    case Train{p, g, NoMom{}, +p1, +p2, +l}:
3509      td0(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p))
3510    case Train{p, g, OneMom{mm}, +p1, +p2, +l}:
3511      td1(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p), Layer.dump(P, X, Y, S, m, mm))
3512    case Train{p, g, TwoMom{mm, vv}, +p1, +p2, +l}:
3513      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))
3514    case Train{p, g, Packed{t}, +p1, +p2, +l}:
3515      tdp(P, g, p1, p2, l, Layer.dump(P, X, Y, S, m, p), pt_cells(t))
3516
3517def f32_at(xs: List<&2, F32>, n: Nat) -> F32:
3518  match xs n:
3519    case Con{x, t} 0n:
3520      x
3521    case Con{x, t} 1n+q:
3522      f32_at(t, q)
3523    case Nil{} _:
3524      0.0
3525
3526def tr0(-P: Type, g: P, r: PA<P>) -> Train<P>:
3527  match r:
3528    case PA{p, xs, off}: Train{p, g, NoMom{}, 0.0, 0.0, 0.0}
3529
3530def tr_m2(-P: Type, p: P, g: P, r: PA<P>) -> Train<P>:
3531  match r:
3532    case PA{mm, xs, off}: Train{p, g, OneMom{mm}, 0.0, 0.0, 0.0}
3533
3534def 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>:
3535  match r:
3536    case PA{p, xs, off}: tr_m2(P, p, g, Layer.load(P, X, Y, S, m, mm, xs, off))
3537
3538def tr_adam2(-P: Type, p: P, mm: P, vv: P, g: P, p1: F32, r: Array<F32> & F32) -> Train<P>:
3539  match r:
3540    case Tuple{xs, p2}: Train{p, g, TwoMom{mm, vv}, p1, p2, 0.0}
3541
3542def tr_adam(-P: Type, p: P, mm: P, vv: P, g: P, r: Array<F32> & F32, +off: U32) -> Train<P>:
3543  match r:
3544    case Tuple{xs, p1}: tr_adam2(P, p, mm, vv, g, p1, Array.get(F32, xs, off))
3545
3546def tr3(-P: Type, p: P, mm: P, g: P, r: PA<P>) -> Train<P>:
3547  match r:
3548    case PA{vv, xs, +off}: tr_adam(P, p, mm, vv, g, Array.get(F32, xs, off), (off + 1 : U32))
3549
3550def 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>:
3551  match r:
3552    case PA{mm, xs, off}:
3553      tr3(P, p, mm, g, Layer.load(P, X, Y, S, m, vv, xs, off))
3554
3555def 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>:
3556  match r:
3557    case PA{p, xs, off}:
3558      tr2(~P, ~X, ~Y, ~S, ~m, p, vv, g, Layer.load(P, X, Y, S, m, mm, xs, off))
3559
3560def trp4(-P: Type, p: P, g: P, t: PTree, +p1: F32, r: Array<F32> & F32) -> Train<P>:
3561  match r:
3562    case Tuple{xs, p2}: Train{p, g, Packed{t}, p1, p2, 0.0}
3563
3564def trp3(-P: Type, p: P, g: P, t: PTree, r: Array<F32> & F32, +off: U32) -> Train<P>:
3565  match r:
3566    case Tuple{xs, p1}: trp4(P, p, g, t, p1, Array.get(F32, xs, off))
3567
3568def trp2(-P: Type, p: P, g: P, r: PTA) -> Train<P>:
3569  match r:
3570    case PTA{t, xs, +off}: trp3(P, p, g, t, Array.get(F32, xs, off), (off + 1 : U32))
3571
3572def trp(-P: Type, g: P, r: PA<P>) -> Train<P>:
3573  match r:
3574    case PA{p, xs, +off}: trp2(P, p, g, pt_load(xs, off))
3575
3576# a training state filled from the numbers Train.dump wrote (its buffers are overwritten in place; packed moments are
3577# built again from the tree's cells)
3578def Train.restore(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, tr: Train<P>, xs: Array<F32>) -> Train<P>:
3579  match tr:
3580    case Train{p, g, NoMom{}, p1, p2, l}:
3581      tr0(P, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3582    case Train{p, g, OneMom{mm}, p1, p2, l}:
3583      tr_m1(~P, ~X, ~Y, ~S, ~m, mm, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3584    case Train{p, g, TwoMom{mm, vv}, p1, p2, l}:
3585      tr1(~P, ~X, ~Y, ~S, ~m, mm, vv, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3586    case Train{p, g, Packed{t}, p1, p2, l}:
3587      trp(P, g, Layer.load(P, X, Y, S, m, p, xs, 0))
3588
3589# ---- writing numbers straight from an array, a chunk at a time ----
3590# File.write_bytes takes a list with one cell per byte; one list for a whole model would take several GB at SmolLM2's
3591# sizes, so each chunk of WCHUNK numbers gets its own list, written and dropped before the next.
3592
3593def WCHUNK() -> U32:
3594  65536
3595
3596# the bytes of xs[i - k + 1 .. i] onto acc, built from the last number back (no reverse)
3597def chunk_bytes(k: Nat, +i: U32, rx: Array<F32> & F32, acc: List<&2, U32>) -> Array<F32> & List<&2, U32>:
3598  match k rx:
3599    case 0n Tuple{a, v}: Tuple{a, acc}
3600    case 1n+q Tuple{a, v}:
3601      +u = bits(v)
3602      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)
3603
3604def wc_ok(-R: Type, rr: File & Result<&1, &1, U32 & String, Unit>, a: Array<F32>, k: File -> Array<F32> -> IO(R)) -> IO(R):
3605  match rr:
3606    case Tuple{f, res}:
3607      do IO<R>:
3608        u : Unit <- IO.pass(Unit, res)
3609        k(f, a)
3610
3611def wc_bytes(-R: Type, f: File, r: Array<F32> & List<&2, U32>, k: File -> Array<F32> -> IO(R)) -> IO(R):
3612  match r:
3613    case Tuple{a, bs}:
3614      do IO<R>:
3615        rr : File & Result<&1, &1, U32 & String, Unit> <- File.write_bytes(f, bs)
3616        wc_ok(R, rr, a, k)
3617
3618# k chunks of a's numbers lo .. n - 1
3619def write_chunks(k: Nat, f: File, a: Array<F32>, +lo: U32, +n: U32) -> IO(File & Array<F32>):
3620  match k:
3621    case 0n:
3622      IO.pure(File & Array<F32>, Tuple{f, a})
3623    case 1n+q:
3624      +hi = U32.min((lo + WCHUNK() : U32), n)
3625      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))
3626
3627def wf_close(r: File & Array<F32>) -> IO(Array<F32>):
3628  match r:
3629    case Tuple{f, a}:
3630      do IO<Array<F32>>:
3631        File.close(f)
3632        return a
3633
3634# cells 0 .. n - 1 of a, as little-endian F32, into a new file at path; a comes back
3635def write_f32(+path: String, a: Array<F32>, +n: U32) -> IO(Array<F32>):
3636  do IO<Array<F32>>:
3637    f : File <- IO.try(File, File.open(path, "w"))
3638    r : File & Array<F32> <- write_chunks(U32.to_nat(U32.div((n + U32.sub(WCHUNK(), 1) : U32), WCHUNK())), f, a, 0, n)
3639    wf_close(r)
3640
3641def ps_fin(-P: Type, +path: String, +n: U32, r: PA<P>) -> IO(P):
3642  match r:
3643    case PA{p, xs, off}:
3644      do IO<P>:
3645        a : Array<F32> <- write_f32(path, xs, n)
3646        return p
3647
3648# every parameter of p, in dump order, into a file at path (Layer.count numbers), through one array
3649def Params.save(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, p: P) -> IO(P):
3650  +n = Layer.count(P, X, Y, S, m)
3651  ps_fin(P, path, n, Layer.store(P, X, Y, S, m, p, Array.new(F32, depth(U32.to_nat(n)), 0.0), 0))
3652
3653# ---- a training state into a file: the parameters, then the moments in use, then Adam's b1^t and b2^t ----
3654
3655def tsa_fin(-P: Type, +path: String, g: P, st: OptState<P>, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3656  match r:
3657    case PA{p, xs, +off}:
3658      do IO<Train<P>>:
3659        a : Array<F32> <- write_f32(path, xs, off)
3660        return Train{p, g, st, p1, p2, l}
3661
3662def tsa_m(-P: Type, +path: String, p: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3663  match r:
3664    case PA{mm, xs, off}: tsa_fin(P, path, g, OneMom{mm}, p1, p2, l, PA{p, xs, off})
3665
3666def tsa_pw(-P: Type, +path: String, p: P, mm: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PA<P>) -> IO(Train<P>):
3667  match r:
3668    case PA{vv, xs, +off}:
3669      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)})
3670
3671def 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>):
3672  match r:
3673    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))
3674
3675def 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>):
3676  match r:
3677    case PA{p, xs, off}: tsa_m(P, path, p, g, p1, p2, l, Layer.store(P, X, Y, S, m, mm, xs, off))
3678
3679def 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>):
3680  match r:
3681    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))
3682
3683def tsp3(-P: Type, +path: String, p: P, g: P, +p1: F32, +p2: F32, +l: F32, r: PTA) -> IO(Train<P>):
3684  match r:
3685    case PTA{t, xs, +off}:
3686      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)})
3687
3688def tsp2(-P: Type, +path: String, g: P, +p1: F32, +p2: F32, +l: F32, t: PTree, r: PA<P>) -> IO(Train<P>):
3689  match r:
3690    case PA{p, xs, +off}: tsp3(P, path, p, g, p1, p2, l, pt_store(t, xs, off))
3691
3692def 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>):
3693  match rt:
3694    case PTU{t, +n}:
3695      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))
3696
3697# the layout Train.dump gives, written from one array instead of a list (the same bytes)
3698def Train.save(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer<P, X, Y, S>, +path: String, tr: Train<P>) -> IO(Train<P>):
3699  match tr:
3700    case Train{p, g, NoMom{}, +p1, +p2, +l}:
3701      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))
3702    case Train{p, g, OneMom{mm}, +p1, +p2, +l}:
3703      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))
3704    case Train{p, g, TwoMom{mm, vv}, +p1, +p2, +l}:
3705      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))
3706    case Train{p, g, Packed{t}, +p1, +p2, +l}:
3707      tsp(~P, ~X, ~Y, ~S, ~m, path, p, g, p1, p2, l, pt_cells(t))