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 the output, the parameters handed back, the saved values 9 # bwd : P -> S -> Y -> P -> Bwd from dY (typed Y) to dX (typed X), adding dP into a gradient of type P 10 # step : Opt -> P -> P -> OptState

-> Step

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` (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 23 import Base 24 import bend-ml-tensor-array@0.1.14.0/main.bend as TA 25 import bend-ml-nat-lemmas@0.1.1.0/main.bend as NL 26 27 # ===================================================================================================== 28 # The layer protocol 29 # ===================================================================================================== 30 31 # what a forward pass returns: the parameters, the output, what the backward pass will need 32 type 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 36 type 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) 40 type 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. 44 def Pair3(A: Type, B: Type, C: Type) -> Type: 45 Two, C> 46 47 def ResidualType(A: Type) -> Type: 48 Pair3(Unit, A, Unit) 49 50 # Owned optimizer buffers: SGD has none, momentum one, Adam/AdamW two. 51 type OptState<-P: Type> is Type: 52 NoMom{} 53 OneMom{m: P} 54 TwoMom{m: P, v: P} 55 56 # Two affine maps are supplied separately: each is called at most once. 57 def OptState.map(-P: Type, -Q: Type, st: OptState

, f: P -> Q, h: P -> Q) -> OptState: 58 match st: 59 case NoMom{}: NoMom{} 60 case OneMom{m}: OneMom{f(m)} 61 case TwoMom{m, v}: TwoMom{f(m), h(v)} 62 63 def OptState.split(-A: Type, -B: Type, st: OptState>) -> Two, OptState>: 64 match st: 65 case NoMom{}: Two{NoMom{}, NoMom{}} 66 case OneMom{Two{a, b}}: Two{OneMom{a}, OneMom{b}} 67 case TwoMom{Two{a, b}, Two{c, d}}: Two{TwoMom{a, c}, TwoMom{b, d}} 68 69 # Pair the buffers present on both sides. Internal callers split one state, then step both children with the same 70 # optimizer, so the variants agree. No buffer is ever allocated to fill an absent side. 71 def OptState.join(-A: Type, -B: Type, a: OptState, b: OptState) -> OptState>: 72 match a b: 73 case OneMom{x} OneMom{y}: OneMom{Two{x, y}} 74 case TwoMom{x, u} TwoMom{y, v}: TwoMom{Two{x, y}, Two{u, v}} 75 case _ _: NoMom{} 76 77 # what an optimizer step returns: the new parameters, the zeroed gradient, the optimizer state 78 type Step<-P: Type> is Type: 79 Step{p: P, g: P, st: OptState

} 80 81 # The optimizer, with the scalars of the current step. Adam's `step` is lr / (1 - b1^t) and `bc2` is 82 # sqrt(1 - b2^t), as PyTorch computes them; `Opt.at` fills them from an `OptCfg` and the step count. 83 # `dec` picks AdamW's decoupled weight decay (w *= 1 - lr*wd) over Adam's (g += wd*w). 84 type Opt is Data: 85 SGD{lr: F32} 86 Momentum{lr: F32, mu: F32} 87 Adam{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, dec: Bool, step: F32, bc2: F32, c1: F32, c2: F32} 88 89 # Parameters and their flat dump in a fixed order. Loading has a separate array cursor (PA). 90 type PL<-P: Type> is Type: 91 PL{p: P, xs: List<&2, F32>} 92 93 # Loading cursor: the shared source array and next element offset, handed back without a list or a clone. 94 # The IO boundary checks the source length against Layer.count before calling the pure loaders. 95 type PA<-P: Type> is Type: 96 PA{p: P, xs: Array, off: U32} 97 98 type Layer<-P: Type, -X: Type, -Y: Type, -S: Type> is Type: 99 Layer{fwd: P -> X -> Fwd, bwd: P -> S -> Y -> P -> Bwd, step: Opt -> P -> P -> OptState

-> Step

, zeros: Unit -> P, dump: P -> PL

, load: P -> Array -> U32 -> PA

, store: P -> Array -> U32 -> PA

, count: Unit -> U32} 100 101 def Layer.fwd(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, x: X) -> Fwd: 102 match l: 103 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 104 fwd(p, x) 105 106 def Layer.bwd(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, s: S, dy: Y, g: P) -> Bwd: 107 match l: 108 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 109 bwd(p, s, dy, g) 110 111 def Layer.step(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, +o: Opt, p: P, g: P, st: OptState

) -> Step

: 112 match l: 113 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 114 step(o, p, g, st) 115 116 def Layer.zeros(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer) -> P: 117 match l: 118 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 119 zeros(Unit{}) 120 121 # every number of the parameters, in the layer's fixed order (the parameters handed back) 122 def Layer.dump(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P) -> PL

: 123 match l: 124 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 125 dump(p) 126 127 # Parameters overwritten from xs at an element offset, in dump order; xs and the next offset come back. 128 def Layer.load(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, xs: Array, off: U32) -> PA

: 129 match l: 130 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: load(p, xs, off) 131 132 # The parameters copied into xs from an element offset, in dump order (load's inverse); xs and the next offset come back. 133 def Layer.store(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, xs: Array, off: U32) -> PA

: 134 match l: 135 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: store(p, xs, off) 136 137 def Layer.count(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer) -> U32: 138 match l: 139 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: count(Unit{}) 140 141 # ===================================================================================================== 142 # Element-wise kernels (templates: one copy per function, no closure at run time) 143 # ===================================================================================================== 144 145 # a[t] = f(a[t]) for the n numbers from t; ra holds the array and a[t], already read 146 def ew1_l(~f: @+a: F32 -> F32, nleft: Nat, +t: U32, ra: Array & F32) -> Array: 147 match nleft ra: 148 case 0n Tuple{a, v}: 149 a 150 case 1n+p Tuple{a, +v}: 151 ew1_l(~f, p, (t + 1 : U32), Array.get(F32, Array.set(F32, a, t, f(v)), (t + 1 : U32))) 152 153 # a[t] = f(a[t], b[t]): the result replaces a; b comes back unchanged 154 def ew2_l(~f: @+a: F32 -> @+b: F32 -> F32, nleft: Nat, +t: U32, ra: Array & F32, rb: Array & F32) -> TA.P2: 155 match nleft ra rb: 156 case 0n Tuple{a, v} Tuple{b, w}: 157 TA.P2{a, b} 158 case 1n+p Tuple{a, +v} Tuple{b, +w}: 159 ew2_l(~f, p, (t + 1 : U32), Array.get(F32, Array.set(F32, a, t, f(v, w)), (t + 1 : U32)), Array.get(F32, b, (t + 1 : U32))) 160 161 def Ew.map(~f: @+a: F32 -> F32, +r: Nat, +c: Nat, m: TA.Mat) -> TA.Mat: 162 match m: 163 case TA.Mat{d}: 164 TA.Mat{ew1_l(~f, Nat.mul(r, c), 0, Array.get(F32, d, 0))} 165 166 def ew2_fin(-r: Nat, -c: Nat, p: TA.P2) -> TA.MPair: 167 match p: 168 case TA.P2{a, b}: 169 TA.MPair{TA.Mat{a}, TA.Mat{b}} 170 171 # (f(a, b), b), element by element 172 def Ew.zip(~f: @+a: F32 -> @+b: F32 -> F32, +r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.MPair: 173 match a b: 174 case TA.Mat{x} TA.Mat{y}: 175 ew2_fin(r, c, ew2_l(~f, Nat.mul(r, c), 0, Array.get(F32, x, 0), Array.get(F32, y, 0))) 176 177 def f32_add(+a: F32, +b: F32) -> F32: 178 (a + b : F32) 179 180 def f32_mul(+a: F32, +b: F32) -> F32: 181 (a * b : F32) 182 183 # (a + b, b) 184 def Mat.add(+r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.MPair: 185 Ew.zip(~f32_add, r, c, a, b) 186 187 def zero_of(+x: F32) -> F32: 188 0.0 189 190 # ===================================================================================================== 191 # Optimizers: one pass over a parameter, its gradient and the two moment buffers; the gradient is zeroed 192 # ===================================================================================================== 193 194 type Q3 is Type: 195 Q3{w: Array, g: Array, m: Array} 196 197 type Q4 is Type: 198 Q4{w: Array, g: Array, m: Array, v: Array} 199 200 type MatStep<-r: Nat, -c: Nat> is Type: 201 MatStep{w: TA.Mat, g: TA.Mat, st: OptState>} 202 203 # w <- w + (-lr)*g, as PyTorch's p.add_(g, alpha=-lr) 204 def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array & F32, rg: Array & F32) -> TA.P2: 205 match nleft rw rg: 206 case 0n Tuple{w, x} Tuple{g, y}: 207 TA.P2{w, g} 208 case 1n+p Tuple{w, +x} Tuple{g, +y}: 209 sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, Array.set(F32, g, t, 0.0), (t + 1 : U32))) 210 211 # buf <- buf*mu + g; w <- w + (-lr)*buf, as PyTorch's SGD with momentum (no dampening, not Nesterov) 212 def mom_l(nleft: Nat, +t: U32, +lr: F32, +mu: F32, rw: Array & F32, rg: Array & F32, rm: Array & F32) -> Q3: 213 match nleft rw rg rm: 214 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, b}: 215 Q3{w, g, m} 216 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}: 217 +nb = ((b * mu : F32) + y : F32) 218 +u = (t + 1 : U32) 219 mom_l(p, u, lr, mu, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * nb : F32) : F32)), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nb), u)) 220 221 # the Adam family, in PyTorch's order of operations (torch.optim.adam._single_tensor_adam): 222 # g' = g + wd*w (Adam) or w' = w*(1 - lr*wd) (AdamW); m = m + (1-b1)*(g' - m) (lerp); v = v*b2 + (1-b2)*g'*g'; 223 # w = w' + (-step*m) / (sqrt(v)/bc2 + eps) 224 def adam_grad(dec: Bool, +wd: F32, +w: F32, +g: F32) -> F32: 225 match dec: 226 case True{}: 227 g 228 case False{}: 229 (g + (w * wd : F32) : F32) 230 231 def adam_decay(dec: Bool, +lr: F32, +wd: F32, +w: F32) -> F32: 232 match dec: 233 case True{}: 234 (w * (1.0 - (lr * wd : F32) : F32) : F32) 235 case False{}: 236 w 237 238 # c1 = 1 - b1 and c2 = 1 - b2 are given as their own constants (0.1, 0.001), as PyTorch computes 1 - beta in float64 239 # before rounding: 1 - b2 taken in F32 from b2 = 0.999 (0.99900001287 in F32) is 1.3e-5 too small, and over 100 steps 240 # that alone put the tiny Llama 21x further from float64 than PyTorch's float32 (NOTES, 2026-10-18) 241 def adam_m(+c1: F32, +m: F32, +g: F32) -> F32: 242 (m + (c1 * (g - m : F32) : F32) : F32) 243 244 def adam_v(+b2: F32, +c2: F32, +v: F32, +g: F32) -> F32: 245 ((v * b2 : F32) + ((c2 * g : F32) * g : F32) : F32) 246 247 def adam_w(+step: F32, +bc2: F32, +eps: F32, +w: F32, +m: F32, +v: F32) -> F32: 248 (w + ((F32.neg(step) * m : F32) / ((F32.sqrt(v) / bc2 : F32) + eps : F32) : F32) : F32) 249 250 def 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, rg: Array & F32, rm: Array & F32, rv: Array & F32) -> Q4: 251 match nleft rw rg rm rv: 252 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}: 253 Q4{w, g, m, v} 254 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}: 255 +gg = adam_grad(dec, wd, x, y) 256 +nm = adam_m(c1, a, gg) 257 +nv = adam_v(b2, c2, b, gg) 258 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), nm, nv) 259 +u = (t + 1 : U32) 260 adam_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nm), u), Array.get(F32, Array.set(F32, v, t, nv), u)) 261 262 # One matrix step, with only the buffers this optimizer reads. 263 def Opt.sgd_fin(-r: Nat, -c: Nat, q: TA.P2) -> MatStep: 264 match q: 265 case TA.P2{w, g}: MatStep{TA.Mat{w}, TA.Mat{g}, NoMom{}} 266 267 def Opt.mom_fin(-r: Nat, -c: Nat, q: Q3) -> MatStep: 268 match q: 269 case Q3{w, g, m}: MatStep{TA.Mat{w}, TA.Mat{g}, OneMom{TA.Mat{m}}} 270 271 def Opt.fin(-r: Nat, -c: Nat, q: Q4) -> MatStep: 272 match q: 273 case Q4{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, TwoMom{TA.Mat{m}, TA.Mat{v}}} 274 275 # A mismatched state is initialized only at the training boundary; Opt.mat requires the matching variant. 276 # The fallback leaves the values untouched, so it cannot fabricate a moment buffer. 277 def Opt.mat(+r: Nat, +c: Nat, +o: Opt, w: TA.Mat, g: TA.Mat, st: OptState>) -> MatStep: 278 match o w g st: 279 case SGD{lr} TA.Mat{wd} TA.Mat{gd} _: 280 Opt.sgd_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, wd, 0), Array.get(F32, gd, 0))) 281 case Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} OneMom{TA.Mat{md}}: 282 Opt.mom_fin(r, c, mom_l(Nat.mul(r, c), 0, lr, mu, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0))) 283 case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} TwoMom{TA.Mat{md}, TA.Mat{vd}}: 284 Opt.fin(r, c, adam_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd2, dec, step, bc2, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0), Array.get(F32, vd, 0))) 285 case _ w g st: 286 MatStep{w, g, st} 287 288 # What the user picks; `Opt.at` turns it into the step's `Opt` from q1 = 1 - b1^t and q2 = 1 - b2^t (t = 1 on the 289 # first step). Adam's c1 and c2 are 1 - b1 and 1 - b2, written as their own constants (adam_m above). 290 type OptCfg is Data: 291 CSGD{lr: F32} 292 CMomentum{lr: F32, mu: F32} 293 CAdam{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32} 294 CAdamW{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32} 295 296 def Opt.at(+cfg: OptCfg, +q1: F32, +q2: F32) -> Opt: 297 match cfg: 298 case CSGD{lr}: 299 SGD{lr} 300 case CMomentum{lr, mu}: 301 Momentum{lr, mu} 302 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 303 Adam{lr, b1, b2, eps, wd, False{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2} 304 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 305 Adam{lr, b1, b2, eps, wd, True{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2} 306 307 def OptCfg.b1(+cfg: OptCfg) -> F32: 308 match cfg: 309 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 310 b1 311 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 312 b1 313 case _: 314 0.0 315 316 def OptCfg.b2(+cfg: OptCfg) -> F32: 317 match cfg: 318 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 319 b2 320 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 321 b2 322 case _: 323 0.0 324 325 def OptCfg.c1(+cfg: OptCfg) -> F32: 326 match cfg: 327 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 328 c1 329 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 330 c1 331 case _: 332 0.0 333 334 def OptCfg.c2(+cfg: OptCfg) -> F32: 335 match cfg: 336 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 337 c2 338 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 339 c2 340 case _: 341 0.0 342 343 # c = 1 - b up to F32's rounding, with 0 < c <= 1 (PyTorch's 0 <= beta < 1). For b in [0, 1), 1 - b taken in F32 is 344 # within half an ulp of the exact value and c's own rounding adds at most as much, so a c written as 1 - b is always 345 # under 2^-23 from it, while a typo (0.01 for 0.001, c1 and c2 swapped) is far above. A NaN fails every comparison. 346 def beta_ok(+b: F32, +c: F32) -> Bool: 347 +d = ((1.0 - b : F32) - c : F32) 348 Bool.and(Bool.and((c > 0.0 : F32), (c <= 1.0 : F32)), ((d * d : F32) <= 0.0000000000000142108547152020037174224853515625 : F32)) 349 350 # True unless an Adam configuration's c1, c2 disagree with its b1, b2. Train.new poisons the state of a configuration 351 # that fails (every loss after the first is NaN); a program checks it first with OptCfg.check to stop with a message. 352 def OptCfg.ok(+cfg: OptCfg) -> Bool: 353 match cfg: 354 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 355 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2)) 356 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 357 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2)) 358 case _: 359 True{} 360 361 def cfg_stop(ok: Bool) -> IO(Unit): 362 match ok: 363 case True{}: 364 IO.pure(Unit, Unit{}) 365 case False{}: 366 IO.die(Unit, 1, "CAdam/CAdamW: c1, c2 do not match b1, b2 (c1 = 1 - b1, c2 = 1 - b2, e.g. 0.9, 0.999 with 0.1, 0.001)") 367 368 def OptCfg.check(+cfg: OptCfg) -> IO(Unit): 369 cfg_stop(OptCfg.ok(cfg)) 370 371 # dump and load for a matrix: its r*c numbers, row by row 372 def mat_dump_fin(-r: Nat, -c: Nat, ml: TA.ML) -> PL>: 373 match ml: 374 case TA.ML{m, xs}: 375 PL{m, xs} 376 377 def Mat.dump(+r: Nat, +c: Nat, m: TA.Mat) -> PL>: 378 mat_dump_fin(r, c, TA.Mat.to_list(r, c, m)) 379 380 def load_cells(left: Nat, +i: U32, +off: U32, rx: Array & F32, dst: Array) -> TA.P2: 381 match left rx: 382 case 0n Tuple{xs, v}: TA.P2{xs, dst} 383 case 1n+q Tuple{xs, v}: 384 load_cells(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, xs, (off + 1 : U32)), Array.set(F32, dst, i, v)) 385 386 def mat_restore_fin(-r: Nat, -c: Nat, +off: U32, q: TA.P2) -> PA>: 387 match q: 388 case TA.P2{xs, dst}: PA{TA.Mat{dst}, xs, off} 389 390 def Mat.restore(+r: Nat, +c: Nat, m: TA.Mat, xs: Array, +off: U32) -> PA>: 391 match m: 392 case TA.Mat{dst}: 393 mat_restore_fin(r, c, (off + U32.from_nat(Nat.mul(r, c)) : U32), load_cells(Nat.mul(r, c), 0, off, Array.get(F32, xs, off), dst)) 394 395 # load_cells' inverse: copies the matrix's cells i.. into xs from `off`; both arrays come back (the matrix's first) 396 def store_cells(left: Nat, +i: U32, +off: U32, rm: Array & F32, xs: Array) -> TA.P2: 397 match left rm: 398 case 0n Tuple{ma, v}: TA.P2{ma, xs} 399 case 1n+q Tuple{ma, v}: 400 store_cells(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, ma, (i + 1 : U32)), Array.set(F32, xs, off, v)) 401 402 def mat_store_fin(-r: Nat, -c: Nat, +off: U32, q: TA.P2) -> PA>: 403 match q: 404 case TA.P2{ma, xs}: PA{TA.Mat{ma}, xs, off} 405 406 def Mat.store(+r: Nat, +c: Nat, m: TA.Mat, xs: Array, +off: U32) -> PA>: 407 match m: 408 case TA.Mat{ma}: 409 mat_store_fin(r, c, (off + U32.from_nat(Nat.mul(r, c)) : U32), store_cells(Nat.mul(r, c), 0, off, Array.get(F32, ma, 0), xs)) 410 411 def pa_second(-A: Type, -B: Type, a: A, rb: PA) -> PA>: 412 match rb: 413 case PA{b, xs, off}: PA{Two{a, b}, xs, off} 414 415 def pl_pair(-A: Type, -B: Type, ra: PL, rb: PL) -> PL>: 416 match ra rb: 417 case PL{a, xa} PL{b, xb}: 418 PL{Two{a, b}, List.append(&2, F32, xa, xb)} 419 420 def pl_second(-A: Type, -B: Type, a: A, rb: PL) -> PL>: 421 match rb: 422 case PL{b, rest}: 423 PL{Two{a, b}, rest} 424 425 # ===================================================================================================== 426 # Linear: y = x·Wᵀ + b, with W stored out x in (PyTorch's nn.Linear layout); saves its input 427 # ===================================================================================================== 428 429 type Lin<-i: Nat, -o: Nat> is Type: 430 Lin{w: TA.Mat, b: TA.Mat<1n, o>} 431 432 def lin_f3(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat, x: TA.Mat, r: TA.MBias) -> Fwd, TA.Mat, TA.Mat>: 433 match r: 434 case TA.MBias{y, b}: 435 Fwd{Lin{w, b}, y, x} 436 437 def lin_f2(+n: Nat, -i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: TA.MMulNT) -> Fwd, TA.Mat, TA.Mat>: 438 match r: 439 case TA.MMulNT{x, w, z}: 440 lin_f3(n, i, o, w, x, TA.Mat.add_row(n, o, z, b)) 441 442 def lin_fwd(+n: Nat, +i: Nat, +o: Nat, par: Nat, p: Lin, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 443 match p: 444 case Lin{w, b}: 445 lin_f2(n, i, o, b, TA.Mat.matmul_nt(n, i, o, par, x, w)) 446 447 def lin_b4(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat, b: TA.Mat<1n, o>, dx: TA.Mat, gw: TA.Mat, r: TA.MColSum) -> Bwd, TA.Mat>: 448 match r: 449 case TA.MColSum{dy, gb}: 450 Bwd{Lin{w, b}, dx, Lin{gw, gb}} 451 452 def lin_b3(+n: Nat, -i: Nat, +o: Nat, w: TA.Mat, b: TA.Mat<1n, o>, dx: TA.Mat, gb: TA.Mat<1n, o>, r: TA.MMulTN) -> Bwd, TA.Mat>: 453 match r: 454 case TA.MMulTN{dy, x, gw}: 455 lin_b4(n, i, o, w, b, dx, gw, TA.Mat.col_sums_acc(n, o, dy, gb)) 456 457 def lin_b2(+n: Nat, +i: Nat, +o: Nat, par: Nat, b: TA.Mat<1n, o>, x: TA.Mat, g: Lin, r: TA.MMul) -> Bwd, TA.Mat>: 458 match g r: 459 case Lin{gw, gb} TA.MMul{dy, w, dx}: 460 lin_b3(n, i, o, w, b, dx, gb, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, gw)) 461 462 # dX = dY·W, dW += dYᵀ·X, db += the column sums of dY; the two sums are added straight into the gradient (the 463 # accumulating products of bend-ml-tensor-array 0.1.10.0), with the bits of computing them apart and adding them 464 def lin_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: Lin, x: TA.Mat, dy: TA.Mat, g: Lin) -> Bwd, TA.Mat>: 465 match p: 466 case Lin{w, b}: 467 lin_b2(n, i, o, par, b, x, g, TA.Mat.matmul_precise(n, o, i, par, dy, w)) 468 469 def pair_s3(-P1: Type, -P2: Type, r1: Step, r2: Step) -> Step>: 470 match r1 r2: 471 case Step{p1, g1, s1} Step{p2, g2, s2}: 472 Step{Two{p1, p2}, Two{g1, g2}, OptState.join(P1, P2, s1, s2)} 473 474 def lin_pair(-i: Nat, -o: Nat, p: Lin) -> Two, TA.Mat<1n, o>>: 475 match p: 476 case Lin{w, b}: Two{w, b} 477 478 def lin_unpair(-i: Nat, -o: Nat, p: Two, TA.Mat<1n, o>>) -> Lin: 479 match p: 480 case Two{w, b}: Lin{w, b} 481 482 def lin_s2(-i: Nat, -o: Nat, a: MatStep, c: MatStep<1n, o>) -> Step>: 483 match a c: 484 case MatStep{w, gw, sw} MatStep{b, gb, sb}: 485 Step{Lin{w, b}, Lin{gw, gb}, OptState.map(Two, TA.Mat<1n, o>>, Lin, OptState.join(TA.Mat, TA.Mat<1n, o>, sw, sb), x => lin_unpair(i, o, x), x => lin_unpair(i, o, x))} 486 487 def lin_step2(+i: Nat, +o: Nat, +op: Opt, p: Lin, g: Lin, st: Two>, OptState>>) -> Step>: 488 match p g st: 489 case Lin{w, b} Lin{gw, gb} Two{sw, sb}: 490 lin_s2(i, o, Opt.mat(o, i, op, w, gw, sw), Opt.mat(1n, o, op, b, gb, sb)) 491 492 def lin_step(+i: Nat, +o: Nat, +op: Opt, p: Lin, g: Lin, st: OptState>) -> Step>: 493 lin_step2(i, o, op, p, g, OptState.split(TA.Mat, TA.Mat<1n, o>, OptState.map(Lin, Two, TA.Mat<1n, o>>, st, x => lin_pair(i, o, x), x => lin_pair(i, o, x)))) 494 495 def lin_dump2(-i: Nat, -o: Nat, r: PL, TA.Mat<1n, o>>>) -> PL>: 496 match r: 497 case PL{Two{w, b}, xs}: 498 PL{Lin{w, b}, xs} 499 500 def lin_dump(+i: Nat, +o: Nat, p: Lin) -> PL>: 501 match p: 502 case Lin{w, b}: 503 lin_dump2(i, o, pl_pair(TA.Mat, TA.Mat<1n, o>, Mat.dump(o, i, w), Mat.dump(1n, o, b))) 504 505 def lin_load3(-i: Nat, -o: Nat, r: PA, TA.Mat<1n, o>>>) -> PA>: 506 match r: 507 case PA{Two{w, b}, xs, off}: PA{Lin{w, b}, xs, off} 508 509 def lin_load2(+i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: PA>) -> PA>: 510 match r: 511 case PA{w, xs, off}: 512 lin_load3(i, o, pa_second(TA.Mat, TA.Mat<1n, o>, w, Mat.restore(1n, o, b, xs, off))) 513 514 def lin_load(+i: Nat, +o: Nat, p: Lin, xs: Array, off: U32) -> PA>: 515 match p: 516 case Lin{w, b}: lin_load2(i, o, b, Mat.restore(o, i, w, xs, off)) 517 518 def lin_store2(+i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: PA>) -> PA>: 519 match r: 520 case PA{w, xs, off}: 521 lin_load3(i, o, pa_second(TA.Mat, TA.Mat<1n, o>, w, Mat.store(1n, o, b, xs, off))) 522 523 def lin_store(+i: Nat, +o: Nat, p: Lin, xs: Array, off: U32) -> PA>: 524 match p: 525 case Lin{w, b}: lin_store2(i, o, b, Mat.store(o, i, w, xs, off)) 526 527 def lin_zeros(+i: Nat, +o: Nat) -> Lin: 528 Lin{TA.Mat.zeros(o, i), TA.Mat.zeros(1n, o)} 529 530 # a linear layer from i to o features on batches of n rows; its products run in 2^par blocks of rows 531 def Linear(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 532 Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => lin_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => lin_zeros(i, o), p => lin_dump(i, o, p), p => xs => off => lin_load(i, o, p, xs, off), p => xs => off => lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))} 533 534 # the first layer of a model, whose input (the data) needs no gradient, as an input without requires_grad in PyTorch: 535 # the backward pass computes dW and db only, skipping dX = dY·W (a third of a linear layer's work), and returns zeros 536 # as dX. Use it only where nothing before the layer has parameters. 537 def linf_b3(+n: Nat, +i: Nat, +o: Nat, w: TA.Mat, b: TA.Mat<1n, o>, gb: TA.Mat<1n, o>, r: TA.MMulTN) -> Bwd, TA.Mat>: 538 match r: 539 case TA.MMulTN{dy, x, gw}: 540 lin_b4(n, i, o, w, b, TA.Mat.zeros(n, i), gw, TA.Mat.col_sums_acc(n, o, dy, gb)) 541 542 def linf_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: Lin, x: TA.Mat, dy: TA.Mat, g: Lin) -> Bwd, TA.Mat>: 543 match p g: 544 case Lin{w, b} Lin{gw, gb}: 545 linf_b3(n, i, o, w, b, gb, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, gw)) 546 547 def LinearFirst(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 548 Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => linf_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => lin_zeros(i, o), p => lin_dump(i, o, p), p => xs => off => lin_load(i, o, p, xs, off), p => xs => off => lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))} 549 550 # ===================================================================================================== 551 # Layers without parameters 552 # ===================================================================================================== 553 554 def pick(b: Bool, +x: F32, +y: F32) -> F32: 555 match b: 556 case True{}: 557 x 558 case False{}: 559 y 560 561 def no_step(-P: Type, +op: Opt, p: P, g: P, st: OptState

) -> Step

: 562 Step{p, g, st} 563 564 # an element-wise activation that saves a copy of its output (relu, sigmoid, tanh: their derivative is a function of 565 # the output). `f` is the activation, `d(dy, y)` the gradient of its input. 566 def act_f2(-n: Nat, -c: Nat, r: TA.MM2) -> Fwd, TA.Mat>: 567 match r: 568 case TA.MM2{y, s}: 569 Fwd{Unit{}, y, s} 570 571 def act_b2(-n: Nat, -c: Nat, r: TA.MPair) -> Bwd>: 572 match r: 573 case TA.MPair{dx, y}: 574 Bwd{Unit{}, dx, Unit{}} 575 576 def ActOut(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 577 Layer{p => x => act_f2(n, c, TA.Mat.clone(n, c, Ew.map(~f, n, c, x))), p => s => dy => g => act_b2(n, c, Ew.zip(~d, n, c, dy, s)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 578 579 # relu(x) = x unless x < 0 (so relu(NaN) = NaN and relu(-0) = -0); its gradient is +0 where the output is <= 0 and 580 # passes everywhere else, NaN included (PyTorch's threshold_backward) 581 def relu_v(+x: F32) -> F32: 582 pick(F32.is_lt(x, 0.0), 0.0, x) 583 584 def relu_d(+dy: F32, +y: F32) -> F32: 585 pick(F32.is_le(y, 0.0), 0.0, dy) 586 587 def ReLU(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 588 ActOut(~relu_v, ~relu_d, n, c) 589 590 # ===================================================================================================== 591 # Seq: la then lb. Parameters, gradients, saved values and optimizer state are pairs of the two layers'. 592 # ===================================================================================================== 593 594 def seq_f3(-P1: Type, -P2: Type, -Z: Type, -S1: Type, -S2: Type, p1: P1, s1: S1, r: Fwd) -> Fwd, Z, Two>: 595 match r: 596 case Fwd{p2, z, s2}: 597 Fwd{Two{p1, p2}, z, Two{s1, s2}} 598 599 def seq_f2(~P2: Type, ~Y: Type, ~Z: Type, ~S2: Type, ~lb: Layer, -P1: Type, -S1: Type, p2: P2, r: Fwd) -> Fwd, Z, Two>: 600 match r: 601 case Fwd{p1, y, s1}: 602 seq_f3(P1, P2, Z, S1, S2, p1, s1, Layer.fwd(P2, Y, Z, S2, lb, p2, y)) 603 604 def seq_fwd(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, x: X) -> Fwd, Z, Two>: 605 match p: 606 case Two{p1, p2}: 607 seq_f2(~P2, ~Y, ~Z, ~S2, ~lb, P1, S1, p2, Layer.fwd(P1, X, Y, S1, la, p1, x)) 608 609 def seq_b3(-P1: Type, -P2: Type, -X: Type, p2: P2, g2: P2, r: Bwd) -> Bwd, X>: 610 match r: 611 case Bwd{p1, dx, g1}: 612 Bwd{Two{p1, p2}, dx, Two{g1, g2}} 613 614 def seq_b2(~P1: Type, ~X: Type, ~Y: Type, ~S1: Type, ~la: Layer, -P2: Type, p1: P1, s1: S1, g1: P1, r: Bwd) -> Bwd, X>: 615 match r: 616 case Bwd{p2, dy, g2}: 617 seq_b3(P1, P2, X, p2, g2, Layer.bwd(P1, X, Y, S1, la, p1, s1, dy, g1)) 618 619 def seq_bwd(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, s: Two, dz: Z, g: Two) -> Bwd, X>: 620 match p s g: 621 case Two{p1, p2} Two{s1, s2} Two{g1, g2}: 622 seq_b2(~P1, ~X, ~Y, ~S1, ~la, P2, p1, s1, g1, Layer.bwd(P2, Y, Z, S2, lb, p2, s2, dz, g2)) 623 624 def seq_s3(-P1: Type, -P2: Type, r1: Step, r2: Step) -> Step>: 625 pair_s3(P1, P2, r1, r2) 626 627 def seq_s2(~P2: Type, ~Y: Type, ~Z: Type, ~S2: Type, ~lb: Layer, -P1: Type, +op: Opt, p2: P2, g2: P2, s2: OptState, r1: Step) -> Step>: 628 seq_s3(P1, P2, r1, Layer.step(P2, Y, Z, S2, lb, op, p2, g2, s2)) 629 630 def seq_step2(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, +op: Opt, p: Two, g: Two, st: Two, OptState>) -> Step>: 631 match p g st: 632 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: 633 seq_s2(~P2, ~Y, ~Z, ~S2, ~lb, P1, op, p2, g2, s2, Layer.step(P1, X, Y, S1, la, op, p1, g1, s1)) 634 635 def seq_step(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, +op: Opt, p: Two, g: Two, st: OptState>) -> Step>: 636 seq_step2(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, op, p, g, OptState.split(P1, P2, st)) 637 638 # dump and load for a pair of layers: the first layer's numbers, then the second's 639 def pair_dump(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two) -> PL>: 640 match p: 641 case Two{p1, p2}: 642 pl_pair(P1, P2, Layer.dump(P1, X1, Y1, S1, la, p1), Layer.dump(P2, X2, Y2, S2, lb, p2)) 643 644 def pair_load2(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer, -P1: Type, p2: P2, r: PA) -> PA>: 645 match r: 646 case PA{p1, xs, off}: 647 pa_second(P1, P2, p1, Layer.load(P2, X2, Y2, S2, lb, p2, xs, off)) 648 649 def pair_load(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, xs: Array, off: U32) -> PA>: 650 match p: 651 case Two{p1, p2}: 652 pair_load2(~P2, ~X2, ~Y2, ~S2, ~lb, P1, p2, Layer.load(P1, X1, Y1, S1, la, p1, xs, off)) 653 654 def pair_store2(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer, -P1: Type, p2: P2, r: PA) -> PA>: 655 match r: 656 case PA{p1, xs, off}: 657 pa_second(P1, P2, p1, Layer.store(P2, X2, Y2, S2, lb, p2, xs, off)) 658 659 def pair_store(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, xs: Array, off: U32) -> PA>: 660 match p: 661 case Two{p1, p2}: 662 pair_store2(~P2, ~X2, ~Y2, ~S2, ~lb, P1, p2, Layer.store(P1, X1, Y1, S1, la, p1, xs, off)) 663 664 def Seq(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer) -> Layer, X, Z, Two>: 665 Layer{p => x => seq_fwd(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, p, x), p => s => dz => g => seq_bwd(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, p, s, dz, g), op => p => g => st => seq_step(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P1, X, Y, S1, la), Layer.zeros(P2, Y, Z, S2, lb)}, p => pair_dump(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X, Y, S1, la) + Layer.count(P2, Y, Z, S2, lb) : U32)} 666 667 # Owned layer composition for dimensions read at run time. Each child record is opened once; each closure goes to 668 # exactly one corresponding closure of the composite layer. The template-based Seq above keeps its existing laws. 669 def seq_at_f2(-P1: Type, -P2: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, p2: P2, fb: P2 -> Y -> Fwd, r: Fwd) -> Fwd, Z, Two>: 670 match r: 671 case Fwd{p1, y, s1}: seq_f3(P1, P2, Z, S1, S2, p1, s1, fb(p2, y)) 672 673 def seq_at_f(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, fa: P1 -> X -> Fwd, fb: P2 -> Y -> Fwd, p: Two, x: X) -> Fwd, Z, Two>: 674 match p: 675 case Two{p1, p2}: seq_at_f2(P1, P2, Y, Z, S1, S2, p2, fb, fa(p1, x)) 676 677 def seq_at_b2(-P1: Type, -P2: Type, -X: Type, -Y: Type, -S1: Type, ba: P1 -> S1 -> Y -> P1 -> Bwd, p1: P1, s1: S1, g1: P1, r: Bwd) -> Bwd, X>: 678 match r: 679 case Bwd{p2, dy, g2}: seq_b3(P1, P2, X, p2, g2, ba(p1, s1, dy, g1)) 680 681 def seq_at_b(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, ba: P1 -> S1 -> Y -> P1 -> Bwd, bb: P2 -> S2 -> Z -> P2 -> Bwd, p: Two, s: Two, dz: Z, g: Two) -> Bwd, X>: 682 match p s g: 683 case Two{p1, p2} Two{s1, s2} Two{g1, g2}: seq_at_b2(P1, P2, X, Y, S1, ba, p1, s1, g1, bb(p2, s2, dz, g2)) 684 685 def seq_at_step(-P1: Type, -P2: Type, sa: Opt -> P1 -> P1 -> OptState -> Step, sb: Opt -> P2 -> P2 -> OptState -> Step, +op: Opt, p: Two, g: Two, st: Two, OptState>) -> Step>: 686 match p g st: 687 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: pair_s3(P1, P2, sa(op, p1, g1, s1), sb(op, p2, g2, s2)) 688 689 def seq_at_dump(-P1: Type, -P2: Type, da: P1 -> PL, db: P2 -> PL, p: Two) -> PL>: 690 match p: 691 case Two{p1, p2}: pl_pair(P1, P2, da(p1), db(p2)) 692 693 def seq_at_load2(-P1: Type, -P2: Type, lb: P2 -> Array -> U32 -> PA, p2: P2, r: PA) -> PA>: 694 match r: 695 case PA{p1, xs, off}: pa_second(P1, P2, p1, lb(p2, xs, off)) 696 697 def seq_at_load(-P1: Type, -P2: Type, la: P1 -> Array -> U32 -> PA, lb: P2 -> Array -> U32 -> PA, p: Two, xs: Array, off: U32) -> PA>: 698 match p: 699 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off)) 700 701 def seq_at_store(-P1: Type, -P2: Type, la: P1 -> Array -> U32 -> PA, lb: P2 -> Array -> U32 -> PA, p: Two, xs: Array, off: U32) -> PA>: 702 match p: 703 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off)) 704 705 def Seq.at(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, a: Layer, b: Layer) -> Layer, X, Z, Two>: 706 match a b: 707 case Layer{fa, ba, sa, za, da, la, ta, na} Layer{fb, bb, sb, zb, db, lb, tb, nb}: 708 Layer{p => x => seq_at_f(P1, P2, X, Y, Z, S1, S2, fa, fb, p, x), p => s => dz => g => seq_at_b(P1, P2, X, Y, Z, S1, S2, ba, bb, p, s, dz, g), op => p => g => st => seq_at_step(P1, P2, sa, sb, op, p, g, OptState.split(P1, P2, st)), u => Two{za(Unit{}), zb(Unit{})}, p => seq_at_dump(P1, P2, da, db, p), p => xs => off => seq_at_load(P1, P2, la, lb, p, xs, off), p => xs => off => seq_at_store(P1, P2, ta, tb, p, xs, off), u => (na(Unit{}) + nb(Unit{}) : U32)} 709 710 # ===================================================================================================== 711 # Losses: the mean loss over the batch, and its gradient with respect to the output 712 # ===================================================================================================== 713 714 type LossOut<-Y: Type> is Type: 715 LossOut{loss: F32, dy: Y} 716 717 # ===================================================================================================== 718 # Training: one step of any layer under any loss 719 # ===================================================================================================== 720 721 # Parameters, gradient (zero between steps), optimizer-specific buffers, Adam's 1 - b1^t and 1 - b2^t, last loss. 722 type Train<-P: Type> is Type: 723 Train{p: P, g: P, st: OptState

, p1: F32, p2: F32, loss: F32} 724 725 def OptState.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, cfg: OptCfg) -> OptState

: 726 match cfg: 727 case CSGD{lr}: NoMom{} 728 case CMomentum{lr, mu}: OneMom{Layer.zeros(P, X, Y, S, m)} 729 case _: TwoMom{Layer.zeros(P, X, Y, S, m), Layer.zeros(P, X, Y, S, m)} 730 731 # 1 - b^t starts at 0; a configuration OptCfg.ok rejects starts at NaN, so every loss after the first is NaN. 732 def q_start(ok: Bool) -> F32: 733 match ok: 734 case True{}: 735 0.0 736 case False{}: 737 (0.0 / 0.0 : F32) 738 739 def Train.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, +cfg: OptCfg, p: P) -> Train

: 740 +q = q_start(OptCfg.ok(cfg)) 741 Train{p, Layer.zeros(P, X, Y, S, m), OptState.new(~P, ~X, ~Y, ~S, ~m, cfg), q, q, 0.0} 742 743 def ts5(-P: Type, +l: F32, +p1: F32, +p2: F32, r: Step

) -> Train

: 744 match r: 745 case Step{p, g, st}: 746 Train{p, g, st, p1, p2, l} 747 748 def ts4(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, +cfg: OptCfg, st: OptState

, +p1: F32, +p2: F32, +l: F32, r: Bwd) -> Train

: 749 match r: 750 case Bwd{p, dx, g}: 751 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 752 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 753 ts5(P, l, q1, q2, Layer.step(P, X, Y, S, m, Opt.at(cfg, q1, q2), p, g, st)) 754 755 def ts3(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, +cfg: OptCfg, p: P, s: S, g: P, st: OptState

, +p1: F32, +p2: F32, r: LossOut) -> Train

: 756 match r: 757 case LossOut{l, dy}: 758 ts4(~P, ~X, ~Y, ~S, ~m, cfg, st, p1, p2, l, Layer.bwd(P, X, Y, S, m, p, s, dy, g)) 759 760 def ts2(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~T: Type, ~m: Layer, ~loss: Y -> T -> LossOut, +cfg: OptCfg, g: P, st: OptState

, +p1: F32, +p2: F32, t: T, r: Fwd) -> Train

: 761 match r: 762 case Fwd{p, y, s}: 763 ts3(~P, ~X, ~Y, ~S, ~m, cfg, p, s, g, st, p1, p2, loss(y, t)) 764 765 # forward, loss, backward and one optimizer step: the same code for every model 766 def Train.step(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~T: Type, ~m: Layer, ~loss: Y -> T -> LossOut, +cfg: OptCfg, tr: Train

, x: X, t: T) -> Train

: 767 match tr: 768 case Train{p, g, st, p1, p2, l}: 769 ts2(~P, ~X, ~Y, ~S, ~T, ~m, ~loss, cfg, g, st, p1, p2, t, Layer.fwd(P, X, Y, S, m, p, x)) 770 771 # One owned model, including dimensions captured at run time. Its fwd, bwd and step closures each run once. 772 def train_at_b(-P: Type, -X: Type, sf: Opt -> P -> P -> OptState

-> Step

, +cfg: OptCfg, st: OptState

, +p1: F32, +p2: F32, +l: F32, r: Bwd) -> Train

: 773 match r: 774 case Bwd{p, dx, g}: 775 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 776 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 777 ts5(P, l, q1, q2, sf(Opt.at(cfg, q1, q2), p, g, st)) 778 779 def train_at_loss(-P: Type, -X: Type, -Y: Type, -S: Type, bf: P -> S -> Y -> P -> Bwd, sf: Opt -> P -> P -> OptState

-> Step

, +cfg: OptCfg, p: P, s: S, g: P, st: OptState

, +p1: F32, +p2: F32, r: LossOut) -> Train

: 780 match r: 781 case LossOut{l, dy}: train_at_b(P, X, sf, cfg, st, p1, p2, l, bf(p, s, dy, g)) 782 783 def train_at_f(-P: Type, -X: Type, -Y: Type, -S: Type, -Target: Type, bf: P -> S -> Y -> P -> Bwd, sf: Opt -> P -> P -> OptState

-> Step

, loss: Y -> Target -> LossOut, +cfg: OptCfg, g: P, st: OptState

, +p1: F32, +p2: F32, t: Target, r: Fwd) -> Train

: 784 match r: 785 case Fwd{p, y, s}: train_at_loss(P, X, Y, S, bf, sf, cfg, p, s, g, st, p1, p2, loss(y, t)) 786 787 def Train.step_at(-P: Type, -X: Type, -Y: Type, -S: Type, -Target: Type, m: Layer, loss: Y -> Target -> LossOut, +cfg: OptCfg, tr: Train

, x: X, t: Target) -> Train

: 788 match m tr: 789 case Layer{ff, bf, sf, zeros, dump, load, store, count} Train{p, g, st, p1, p2, l}: 790 train_at_f(P, X, Y, S, Target, bf, sf, loss, cfg, g, st, p1, p2, t, ff(p, x)) 791 792 # ===================================================================================================== 793 # Reading and printing numbers (bend-ml's and llama.bend's loaders, MIT) 794 # ===================================================================================================== 795 796 # U32 and F32 hold the same Word(32n): these are bitcasts 797 def bitcast(u: U32) -> F32: 798 match u: 799 case U32{w}: 800 F32{w} 801 802 def bits(x: F32) -> U32: 803 match x: 804 case F32{w}: 805 U32{w} 806 807 def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32: 808 bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32)) 809 810 def f32s(bs: List<&2, U32>) -> List<&2, F32>: 811 match bs: 812 case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}: 813 f32_le(b0, b1, b2, b3) <> f32s(t) 814 case _: 815 Nil{} 816 817 def rd_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>): 818 match m: 819 case Tuple{f, r}: 820 do IO>: 821 bs : List<&2, U32> <- IO.pass(List<&2, U32>, r) 822 File.close(f) 823 return bs 824 825 # n bytes from offset off 826 def read_at(path: String, +off: U32, +n: U32) -> IO(List<&2, U32>): 827 do IO>: 828 f : File <- IO.try(File, File.open(path, "r")) 829 m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, off, n) 830 rd_done(m) 831 832 def mat_or_die(-r: Nat, -c: Nat, +path: String, m: Maybe<&1, TA.Mat>) -> IO(TA.Mat): 833 match m: 834 case None{}: 835 IO.die(TA.Mat, 1, path ++ ": wrong number of values for this shape") 836 case Some{x}: 837 IO.pure(TA.Mat, x) 838 839 # an r x c matrix stored as r*c little-endian F32 from byte off of a file 840 def Mat.load(+r: Nat, +c: Nat, +path: String, +off: U32) -> IO(TA.Mat): 841 do IO>: 842 bs : List<&2, U32> <- read_at(path, off, (U32.from_nat(Nat.mul(r, c)) * 4 : U32)) 843 mat_or_die(r, c, path, TA.Mat.of_list(r, c, f32s(bs))) 844 845 def words(xs: List<&2, F32>) -> String: 846 match xs: 847 case Nil{}: 848 "" 849 case Con{h, t}: 850 " " ++ U32.show(bits(h)) ++ words(t) 851 852 def ml_words(-r: Nat, -c: Nat, m: TA.ML) -> String: 853 match m: 854 case TA.ML{d, xs}: 855 words(xs) 856 857 # one line: the name, then each number as the decimal of its 32 bits (exact, for the tests in reference/) 858 def Mat.print(+r: Nat, +c: Nat, +name: String, m: TA.Mat) -> IO(Unit): 859 IO.print(name ++ ml_words(r, c, TA.Mat.to_list(r, c, m))) 860 861 def F32.print(+name: String, x: F32) -> IO(Unit): 862 IO.print(name ++ " " ++ U32.show(bits(x))) 863 864 # the smallest d with 2^d >= n (TA.cap_depth's answer), by doubling a U32: log2(n) cheap steps, where cap_depth 865 # compares Nats one unit at a time at each doubling; f is fuel (32 doublings cover every U32) 866 def depth_go(f: Nat, ok: Bool, +n: U32, +p: U32, d: Nat) -> Nat: 867 match f ok: 868 case _ True{}: 869 d 870 case 0n False{}: 871 d 872 case 1n+g False{}: 873 depth_go(g, U32.is_le(n, (p + p : U32)), n, (p + p : U32), 1n+d) 874 875 def depth(+n: Nat) -> Nat: 876 depth_go(32n, U32.is_le(U32.from_nat(n), 1), U32.from_nat(n), 1, 0n) 877 878 # n zeros, in an array with room for them 879 def zeros_a(+n: Nat) -> Array: 880 Array.new(F32, depth(n), 0.0) 881 882 # ===================================================================================================== 883 # More element-wise activations 884 # ===================================================================================================== 885 886 # an activation whose derivative needs its input (silu, gelu): saves a copy of the input 887 def acti_f3(~f: @+a: F32 -> F32, +n: Nat, +c: Nat, m: TA.MM2) -> Fwd, TA.Mat>: 888 match m: 889 case TA.MM2{x, s}: 890 Fwd{Unit{}, Ew.map(~f, n, c, x), s} 891 892 def ActIn(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 893 Layer{p => x => acti_f3(~f, n, c, TA.Mat.clone(n, c, x)), p => s => dy => g => act_b2(n, c, Ew.zip(~d, n, c, dy, s)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 894 895 def sigmoid_v(+x: F32) -> F32: 896 (1.0 / (1.0 + F32.exp(F32.neg(x)) : F32) : F32) 897 898 # PyTorch's sigmoid_backward: dy * (1 - y) * y 899 def sigmoid_d(+dy: F32, +y: F32) -> F32: 900 ((dy * (1.0 - y : F32) : F32) * y : F32) 901 902 def Sigmoid(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 903 ActOut(~sigmoid_v, ~sigmoid_d, n, c) 904 905 def tanh_v(+x: F32) -> F32: 906 F32.tanh(x) 907 908 def tanh_d(+dy: F32, +y: F32) -> F32: 909 (dy * (1.0 - (y * y : F32) : F32) : F32) 910 911 def Tanh(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 912 ActOut(~tanh_v, ~tanh_d, n, c) 913 914 # silu(x) = x * sigmoid(x), computed as PyTorch's CPU kernel: x / (1 + exp(-x)); 915 # its gradient dy * s * (1 + x * (1 - s)) with s = sigmoid(x) 916 def silu_v(+x: F32) -> F32: 917 (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32) 918 919 def silu_d(+dy: F32, +x: F32) -> F32: 920 +s = sigmoid_v(x) 921 ((dy * s : F32) * (1.0 + (x * (1.0 - s : F32) : F32) : F32) : F32) 922 923 def SiLU(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 924 ActIn(~silu_v, ~silu_d, n, c) 925 926 # GELU in its tanh form (PyTorch's approximate="tanh"): 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3))). 927 # The exact form, with erf, is GELUErf below. 928 def gelu_inner(+x: F32) -> F32: 929 (0.7978845608028654 * (x + (0.044715 * ((x * x : F32) * x : F32) : F32) : F32) : F32) 930 931 def gelu_v(+x: F32) -> F32: 932 (((0.5 * x : F32) * (1.0 + F32.tanh(gelu_inner(x)) : F32)) : F32) 933 934 def gelu_d(+dy: F32, +x: F32) -> F32: 935 +th = F32.tanh(gelu_inner(x)) 936 +left = (0.5 * x : F32) 937 +dleft = (0.5 * (1.0 + th : F32) : F32) 938 +dright = ((left * (1.0 - (th * th : F32) : F32) : F32) * (0.7978845608028654 * (1.0 + ((3.0 * 0.044715 : F32) * (x * x : F32) : F32) : F32) : F32) : F32) 939 (dy * (dleft + dright : F32) : F32) 940 941 def GELU(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 942 ActIn(~gelu_v, ~gelu_d, n, c) 943 944 # ===================================================================================================== 945 # erf in plain Bend, and GELU in its exact form 946 # ===================================================================================================== 947 948 # erf(x) = 2/sqrt(pi) times the integral of exp(-t^2) from 0 to x. Base has no erf, so it is written here with F32 949 # additions, multiplications and comparisons only: no libm call, so every target runs the same operations (the emitted C 950 # turns fused multiply-adds off); the C and JavaScript programs give the same bits (CUDA is not checked yet). erf is 951 # odd, erf(-x) = -erf(x), so the work is done on a = |x| and the sign put back, in three kinds of pieces: 952 # - a < 0.84375: erf(x) = x + x P(x^2), with P of degree 6 (fdlibm's erff has this x + x y shape there); it takes 953 # x itself, so erf(-0) = -0, and NaN reaches it (every comparison with NaN is false) and comes out NaN; 954 # - 0.84375 <= a < 3.92, in six intervals around c = 1, 1.5, ..., 3.5: erf(a) = K + S(a - c), where K is erf(c) 955 # rounded to F32 and S a polynomial of degree 7 or 8 whose value is small next to K, so the last addition rounds 956 # once and carries almost all of the error (a - c is exact: a and c are within a factor of 2); 957 # - a >= 3.92, infinity included: 1, since erf(a) rounds to 1 in F32 from a = 3.9192 on. 958 # The coefficients are least-squares fits against a 50-digit erf on Chebyshev points, rounded to F32. Measured on all 959 # 2^32 inputs against glibc's double erf: at most 0.93 ulp off, the correctly rounded F32 on 97.1% of the inputs, and 1 960 # ulp from it on the rest; NaN for every NaN. 961 962 # one Horner step, c + t s; erf_hm is the step for a negative coefficient -c: t s - c gives the bits of -c + t s 963 def erf_hp(+c: F32, +t: F32, +s: F32) -> F32: 964 (c + (t * s : F32) : F32) 965 966 def erf_hm(+c: F32, +t: F32, +s: F32) -> F32: 967 ((t * s : F32) - c : F32) 968 969 def erf_small(+x: F32) -> F32: 970 +z = (x * x : F32) 971 (x + (x * erf_hp(0.12837916612625122, z, erf_hm(0.37612634897232056, z, erf_hp(0.11283732205629349, z, erf_hm(0.026861993595957756, z, erf_hp(0.005208558402955532, z, erf_hm(0.000823960464913398, z, 8.829222497297451e-05)))))) : F32) : F32) 972 973 # [0.84375, 1.25): t = a - 1.0 974 def erf_p0(+t: F32) -> F32: 975 (0.8427007794380188 + erf_hp(1.3508899776581984e-08, t, erf_hp(0.41510748863220215, t, erf_hm(0.41510751843452454, t, erf_hp(0.13837093114852905, t, erf_hp(0.06918490678071976, t, erf_hm(0.06929104775190353, t, erf_hp(0.004778923466801643, t, erf_hp(0.017144912853837013, t, F32.neg(0.010648881085216999))))))))) : F32) 976 977 # [1.25, 1.75): t = a - 1.5 978 def erf_p1(+t: F32) -> F32: 979 (0.9661051630973816 + erf_hm(1.6622220044837377e-08, t, erf_hp(0.11893028765916824, t, erf_hm(0.17839543521404266, t, erf_hp(0.1387522667646408, t, erf_hm(0.04459875077009201, t, erf_hm(0.014879402704536915, t, erf_hp(0.019323812797665596, t, erf_hm(0.004488371312618256, t, F32.neg(0.0023522074334323406))))))))) : F32) 980 981 # [1.75, 2.25): t = a - 2.0 982 def erf_p2(+t: F32) -> F32: 983 (0.9953222870826721 + erf_hm(2.2064206461891445e-08, t, erf_hp(0.02066698484122753, t, erf_hm(0.04133396968245506, t, erf_hp(0.04822301119565964, t, erf_hm(0.03444507718086243, t, erf_hp(0.013088003732264042, t, erf_hp(0.00046277992078103125, t, erf_hm(0.0033736424520611763, t, 0.0015402256976813078)))))))) : F32) 984 985 # [2.25, 2.75): t = a - 2.5 986 def erf_p3(+t: F32) -> F32: 987 (0.9995930194854736 + erf_hp(2.846470792405853e-08, t, erf_hp(0.002178284339606762, t, erf_hm(0.005445694085210562, t, erf_hp(0.008350063115358353, t, erf_hm(0.008623690344393253, t, erf_hp(0.00611893879249692, t, erf_hm(0.002765072975307703, t, 0.0005077228415757418))))))) : F32) 988 989 # [2.75, 3.25): t = a - 3.0 990 def erf_p4(+t: F32) -> F32: 991 (0.9999778866767883 + erf_hp(2.2842012725732275e-08, t, erf_hp(0.00013925303937867284, t, erf_hm(0.0004177672672085464, t, erf_hp(0.0007891026325523853, t, erf_hm(0.0010437482269480824, t, erf_hp(0.001016448251903057, t, erf_hm(0.0007547215791419148, t, 0.0003925197524949908))))))) : F32) 992 993 # [3.25, 3.92): t = a - 3.5 994 def erf_p5(+t: F32) -> F32: 995 (0.9999992847442627 + erf_hm(2.7855115547481546e-08, t, erf_hp(5.400151167123113e-06, t, erf_hm(1.8894650565925986e-05, t, erf_hp(4.221926792524755e-05, t, erf_hm(6.778862734790891e-05, t, erf_hp(8.422814426012337e-05, t, erf_hm(7.95367595856078e-05, t, 4.151461325818673e-05))))))) : F32) 996 997 def erf_sign(+x: F32, +e: F32) -> F32: 998 pick(F32.is_lt(x, 0.0), F32.neg(e), e) 999 1000 # the pieces are chosen from the top down, so that only one polynomial is computed (each test is a match on a Bool 1001 # parameter); b is "a is at least the piece's lower end" 1002 def erf_i0(b: Bool, +x: F32, +a: F32) -> F32: 1003 match b: 1004 case True{}: erf_sign(x, erf_p0((a - 1.0 : F32))) 1005 case False{}: erf_small(x) 1006 1007 def erf_i1(b: Bool, +x: F32, +a: F32) -> F32: 1008 match b: 1009 case True{}: erf_sign(x, erf_p1((a - 1.5 : F32))) 1010 case False{}: erf_i0(F32.is_ge(a, 0.84375), x, a) 1011 1012 def erf_i2(b: Bool, +x: F32, +a: F32) -> F32: 1013 match b: 1014 case True{}: erf_sign(x, erf_p2((a - 2.0 : F32))) 1015 case False{}: erf_i1(F32.is_ge(a, 1.25), x, a) 1016 1017 def erf_i3(b: Bool, +x: F32, +a: F32) -> F32: 1018 match b: 1019 case True{}: erf_sign(x, erf_p3((a - 2.5 : F32))) 1020 case False{}: erf_i2(F32.is_ge(a, 1.75), x, a) 1021 1022 def erf_i4(b: Bool, +x: F32, +a: F32) -> F32: 1023 match b: 1024 case True{}: erf_sign(x, erf_p4((a - 3.0 : F32))) 1025 case False{}: erf_i3(F32.is_ge(a, 2.25), x, a) 1026 1027 def erf_i5(b: Bool, +x: F32, +a: F32) -> F32: 1028 match b: 1029 case True{}: erf_sign(x, erf_p5((a - 3.5 : F32))) 1030 case False{}: erf_i4(F32.is_ge(a, 2.75), x, a) 1031 1032 def erf_top(b: Bool, +x: F32, +a: F32) -> F32: 1033 match b: 1034 case True{}: erf_sign(x, 1.0) 1035 case False{}: erf_i5(F32.is_ge(a, 3.25), x, a) 1036 1037 def erf(+x: F32) -> F32: 1038 +a = F32.abs(x) 1039 erf_top(F32.is_ge(a, 3.92), x, a) 1040 1041 # GELU in its exact form, PyTorch's default (approximate="none"): gelu(x) = x/2 (1 + erf(x / sqrt 2)), in PyTorch's CPU 1042 # order, (x * 0.5) * (1 + erf(x * 0.70710677)). Its gradient is cdf + x pdf, with cdf = 0.5 (1 + erf(x / sqrt 2)) and 1043 # pdf = exp(-x^2 / 2) / sqrt(2 pi), computed as PyTorch does: 0.5 (1 + erf(x * 0.70710677)) and 1044 # 0.3989423 * exp((x * x) * -0.5). The exponential is F32.exp, as in SiLU (so on a GPU the gradient would take CUDA's 1045 # float exp). 1046 def gelu_erf_cdf(+x: F32) -> F32: 1047 (0.5 * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32) 1048 1049 def gelu_erf_v(+x: F32) -> F32: 1050 ((x * 0.5 : F32) * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32) 1051 1052 def gelu_erf_d(+dy: F32, +x: F32) -> F32: 1053 +pdf = (0.3989422804014327 * F32.exp(((x * x : F32) * F32.neg(0.5) : F32)) : F32) 1054 (dy * (gelu_erf_cdf(x) + (x * pdf : F32) : F32) : F32) 1055 1056 def GELUErf(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1057 ActIn(~gelu_erf_v, ~gelu_erf_d, n, c) 1058 1059 # ===================================================================================================== 1060 # Dropout: each number is kept with probability 1 - p and scaled by 1 / (1 - p), or set to 0. The mask comes from a 1061 # counter-based hash of (seed, step, index), so a run is reproducible and a twin can compute the same mask. The step 1062 # counter is the layer's "parameter": the optimizer step only increments it, and checkpoints keep it. 1063 # ===================================================================================================== 1064 1065 type DropP is Data: 1066 DropP{t: U32} 1067 1068 # lowbias32 (Chris Wellons): an invertible mix of the 32 bits 1069 def mix32(+x: U32) -> U32: 1070 +a = U32.xor(x, U32.shrn(x, 16n)) 1071 +b = (a * 2146121005 : U32) 1072 +c = U32.xor(b, U32.shrn(b, 15n)) 1073 +d = (c * 2221713035 : U32) 1074 U32.xor(d, U32.shrn(d, 16n)) 1075 1076 # the factor of number i: 1 / (1 - p) if its uniform u = (h >> 8) / 2^24 is >= p, else 0 1077 def drop_k(+key: U32, +i: U32, +p: F32, +scale: F32) -> F32: 1078 pick(F32.is_lt((U32.to_f32(U32.shrn(mix32((key + i : U32)), 8n)) / 16777216.0 : F32), p), 0.0, scale) 1079 1080 def drop_l(nleft: Nat, +i: U32, +key: U32, +p: F32, +scale: F32, a: Array) -> Array: 1081 match nleft: 1082 case 0n: 1083 a 1084 case 1n+q: 1085 drop_l(q, (i + 1 : U32), key, p, scale, Array.set(F32, a, i, drop_k(key, i, p, scale))) 1086 1087 def drop_mask(+n: Nat, +c: Nat, +seed: U32, +t: U32, +p: F32) -> TA.Mat: 1088 TA.Mat{drop_l(Nat.mul(n, c), 0, mix32(U32.xor(seed, mix32(t))), p, (1.0 / (1.0 - p : F32) : F32), zeros_a(Nat.mul(n, c)))} 1089 1090 def drop_f2(-n: Nat, -c: Nat, +t: U32, r: TA.MPair) -> Fwd, TA.Mat>: 1091 match r: 1092 case TA.MPair{y, k}: 1093 Fwd{DropP{t}, y, k} 1094 1095 def drop_fwd(+n: Nat, +c: Nat, +seed: U32, +p: F32, d: DropP, x: TA.Mat) -> Fwd, TA.Mat>: 1096 match d: 1097 case DropP{+t}: 1098 drop_f2(n, c, t, Ew.zip(~f32_mul, n, c, x, drop_mask(n, c, seed, t, p))) 1099 1100 def drop_b2(-n: Nat, -c: Nat, d: DropP, g: DropP, r: TA.MPair) -> Bwd>: 1101 match r: 1102 case TA.MPair{dx, k}: 1103 Bwd{d, dx, g} 1104 1105 def drop_step(+op: Opt, d: DropP, g: DropP, st: OptState) -> Step: 1106 match d: 1107 case DropP{t}: 1108 Step{DropP{(t + 1 : U32)}, g, st} 1109 1110 def drop_dump(d: DropP) -> PL: 1111 match d: 1112 case DropP{+t}: 1113 PL{DropP{t}, [bitcast(t)]} 1114 1115 def drop_load_fin(+off: U32, r: Array & F32) -> PA: 1116 match r: 1117 case Tuple{a, v}: PA{DropP{bits(v)}, a, (off + 1 : U32)} 1118 1119 def drop_load(d: DropP, xs: Array, +off: U32) -> PA: 1120 drop_load_fin(off, Array.get(F32, xs, off)) 1121 1122 # the counter, bitcast to one F32 as dump writes it 1123 def drop_store(d: DropP, xs: Array, +off: U32) -> PA: 1124 match d: 1125 case DropP{+c}: PA{DropP{c}, Array.set(F32, xs, off, bitcast(c)), (off + 1 : U32)} 1126 1127 def Dropout(+n: Nat, +c: Nat, +p: F32, +seed: U32) -> Layer, TA.Mat, TA.Mat>: 1128 Layer{d => x => drop_fwd(n, c, seed, p, d, x), d => s => dy => g => drop_b2(n, c, d, g, Ew.zip(~f32_mul, n, c, dy, s)), op => d => g => st => drop_step(op, d, g, st), u => DropP{0}, d => drop_dump(d), d => xs => off => drop_load(d, xs, off), d => xs => off => drop_store(d, xs, off), u => 1} 1129 1130 # ===================================================================================================== 1131 # Row kernels over arrays: loops over a row of a flat array by index. Every sum runs left to right from +0.0 and every 1132 # expression keeps its order of operations, the ones the row-wise layers had when they worked on lists of rows, so 1133 # they give the same bits (NOTES.md, M6) 1134 # ===================================================================================================== 1135 1136 type D2 is Type: 1137 D2{a: Array, b: Array, s: F32} 1138 1139 type A2 is Type: 1140 A2{a: Array, b: Array} 1141 1142 # a[ia]·b[ib] + a[ia + sa]·b[ib + sb] + ... (k terms) added to acc left to right; ra and rb hold the 1143 # first pair, already read 1144 def sdot(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array & F32, rb: Array & F32) -> D2: 1145 match k ra rb: 1146 case 0n Tuple{a, x} Tuple{b, y}: 1147 D2{a, b, acc} 1148 case 1n+p Tuple{a, +x} Tuple{b, +y}: 1149 sdot(p, (ia + sa : U32), (ib + sb : U32), sa, sb, (acc + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32))) 1150 1151 def dot0(+k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a: Array, b: Array) -> D2: 1152 sdot(k, ia, ib, sa, sb, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib)) 1153 1154 1155 type AS is Type: 1156 AS{a: Array, s: F32} 1157 1158 # the largest of a[idx] (held in r) and the k numbers after it: max(max(a0, a1), a2) ..., from the first 1159 def pmax(k: Nat, +idx: U32, cur: F32, r: Array & F32) -> AS: 1160 match k r: 1161 case 0n Tuple{a, +x}: 1162 AS{a, F32.max(cur, x)} 1163 case 1n+q Tuple{a, +x}: 1164 pmax(q, (idx + 1 : U32), F32.max(cur, x), Array.get(F32, a, (idx + 1 : U32))) 1165 1166 def pmax0(k: Nat, +idx: U32, r: Array & F32) -> AS: 1167 match k r: 1168 case 0n Tuple{a, +x}: 1169 AS{a, x} 1170 case 1n+q Tuple{a, +x}: 1171 pmax(q, (idx + 1 : U32), x, Array.get(F32, a, (idx + 1 : U32))) 1172 1173 # a[j] = exp(a[j] - m) for j = idx..idx + k, their sum added to acc in order 1174 def pexp(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array & F32) -> AS: 1175 match k r: 1176 case 0n Tuple{a, +x}: 1177 +e = F32.exp((x - m : F32)) 1178 AS{Array.set(F32, a, idx, e), (acc + e : F32)} 1179 case 1n+q Tuple{a, +x}: 1180 +e = F32.exp((x - m : F32)) 1181 pexp(q, (idx + 1 : U32), m, (acc + e : F32), Array.get(F32, Array.set(F32, a, idx, e), (idx + 1 : U32))) 1182 1183 # the array and how many numbers the last loop of a softmax divided 1184 type PC is Type: 1185 PC{a: Array, c: Nat} 1186 1187 def pc_a(r: PC) -> Array: 1188 match r: 1189 case PC{a, c}: 1190 a 1191 1192 # a[j] = a[j] / s for j = idx..idx + k, counted onto c 1193 def pdiv(k: Nat, +idx: U32, +s: F32, c: Nat, r: Array & F32) -> PC: 1194 match k r: 1195 case 0n Tuple{a, +x}: 1196 PC{Array.set(F32, a, idx, (x / s : F32)), 1n+c} 1197 case 1n+q Tuple{a, +x}: 1198 pdiv(q, (idx + 1 : U32), s, 1n+c, Array.get(F32, Array.set(F32, a, idx, (x / s : F32)), (idx + 1 : U32))) 1199 1200 def psm3(k: Nat, +off: U32, r: AS) -> PC: 1201 match r: 1202 case AS{a, +s}: 1203 pdiv(k, off, s, 0n, Array.get(F32, a, off)) 1204 1205 def psm2(+k: Nat, +off: U32, r: AS) -> PC: 1206 match r: 1207 case AS{a, +m}: 1208 psm3(k, off, pexp(k, off, m, 0.0, Array.get(F32, a, off))) 1209 1210 # softmax in place over the n numbers of p from off (n >= 1), as PyTorch: exp(x - max) / sum; with the count of the 1211 # numbers it divided (the law causal reads it) 1212 def psmc(+n: Nat, +off: U32, p: Array) -> PC: 1213 psm2(Nat.sub(n, 1n), off, pmax0(Nat.sub(n, 1n), off, Array.get(F32, p, off))) 1214 1215 def psm(+n: Nat, +off: U32, p: Array) -> Array: 1216 pc_a(psmc(n, off, p)) 1217 1218 # dst[di..] = src[si..], cnt + 1 numbers; r holds src[si], already read 1219 def copy_n(cnt: Nat, +si: U32, +di: U32, o: Array, r: Array & F32) -> A2: 1220 match cnt r: 1221 case 0n Tuple{s, +x}: 1222 A2{s, Array.set(F32, o, di, x)} 1223 case 1n+n Tuple{s, +x}: 1224 copy_n(n, (si + 1 : U32), (di + 1 : U32), Array.set(F32, o, di, x), Array.get(F32, s, (si + 1 : U32))) 1225 1226 # ===================================================================================================== 1227 # Row-wise layers: softmax, log-softmax, RMSNorm, LayerNorm (each row of an n x c matrix on its own) 1228 # ===================================================================================================== 1229 1230 # rows i.. of x (c wide, left more after i), each softmaxed in place (psm) 1231 def smf_rows(left: Nat, +i: U32, +c: Nat, x: Array) -> Array: 1232 match left: 1233 case 0n: 1234 psm(c, (i * U32.from_nat(c) : U32), x) 1235 case 1n+q: 1236 smf_rows(q, (i + 1 : U32), c, psm(c, (i * U32.from_nat(c) : U32), x)) 1237 1238 # the output, and a copy of it saved for the backward pass 1239 def out_saved(-n: Nat, -c: Nat, r: TA.MM2) -> Fwd, TA.Mat>: 1240 match r: 1241 case TA.MM2{a, b}: 1242 Fwd{Unit{}, a, b} 1243 1244 def sm_fwd(+n: Nat, +c: Nat, x: TA.Mat) -> Fwd, TA.Mat>: 1245 match x: 1246 case TA.Mat{a}: 1247 out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{smf_rows(Nat.sub(n, 1n), 0, c, a)})) 1248 1249 # dy[j] = y[j] * (dy[j] - s), j = idx..idx + k; rd and ry hold dy[idx] and y[idx] 1250 def smb_fix(k: Nat, +idx: U32, +s: F32, rd: Array & F32, ry: Array & F32) -> A2: 1251 match k rd ry: 1252 case 0n Tuple{d, +g} Tuple{y, +v}: 1253 A2{Array.set(F32, d, idx, (v * (g - s : F32) : F32)), y} 1254 case 1n+q Tuple{d, +g} Tuple{y, +v}: 1255 smb_fix(q, (idx + 1 : U32), s, Array.get(F32, Array.set(F32, d, idx, (v * (g - s : F32) : F32)), (idx + 1 : U32)), Array.get(F32, y, (idx + 1 : U32))) 1256 1257 def smb_row(+c: Nat, +off: U32, r: D2) -> A2: 1258 match r: 1259 case D2{d, y, +s}: 1260 smb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off)) 1261 1262 # softmax backward, as PyTorch's: dx = y * (dy - sum(dy * y)), row after row, in place over dy 1263 def smb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2: 1264 match left st: 1265 case 0n A2{d, y}: 1266 smb_row(c, (i * U32.from_nat(c) : U32), dot0(c, (i * U32.from_nat(c) : U32), (i * U32.from_nat(c) : U32), 1, 1, d, y)) 1267 case 1n+q A2{d, y}: 1268 smb_rows(q, (i + 1 : U32), c, smb_row(c, (i * U32.from_nat(c) : U32), dot0(c, (i * U32.from_nat(c) : U32), (i * U32.from_nat(c) : U32), 1, 1, d, y))) 1269 1270 def dx_fin(-n: Nat, -c: Nat, r: A2) -> Bwd>: 1271 match r: 1272 case A2{d, y}: 1273 Bwd{Unit{}, TA.Mat{d}, Unit{}} 1274 1275 def sm_bwd(+n: Nat, +c: Nat, s: TA.Mat, dy: TA.Mat) -> Bwd>: 1276 match s dy: 1277 case TA.Mat{y} TA.Mat{d}: 1278 dx_fin(n, c, smb_rows(Nat.sub(n, 1n), 0, c, A2{d, y})) 1279 1280 # softmax over each row; saves its output 1281 def Softmax(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1282 Layer{p => x => sm_fwd(n, c, x), p => s => dy => g => sm_bwd(n, c, s, dy), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1283 1284 # log_softmax(x) = (x - max) - log(sum(exp(x - max))), as PyTorch's CPU kernel 1285 1286 # acc + exp(a[j] - m) for j = idx..idx + k, in order 1287 def lsum(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array & F32) -> AS: 1288 match k r: 1289 case 0n Tuple{a, +x}: 1290 AS{a, (acc + F32.exp((x - m : F32)) : F32)} 1291 case 1n+q Tuple{a, +x}: 1292 lsum(q, (idx + 1 : U32), m, (acc + F32.exp((x - m : F32)) : F32), Array.get(F32, a, (idx + 1 : U32))) 1293 1294 # a[j] = (a[j] - m) - l for j = idx..idx + k 1295 def lw(k: Nat, +idx: U32, +m: F32, +l: F32, r: Array & F32) -> Array: 1296 match k r: 1297 case 0n Tuple{a, +x}: 1298 Array.set(F32, a, idx, ((x - m : F32) - l : F32)) 1299 case 1n+q Tuple{a, +x}: 1300 lw(q, (idx + 1 : U32), m, l, Array.get(F32, Array.set(F32, a, idx, ((x - m : F32) - l : F32)), (idx + 1 : U32))) 1301 1302 def lsm_r3(+c: Nat, +off: U32, +m: F32, r: AS) -> Array: 1303 match r: 1304 case AS{a, +s}: 1305 lw(Nat.sub(c, 1n), off, m, F32.log(s), Array.get(F32, a, off)) 1306 1307 def lsm_r2(+c: Nat, +off: U32, r: AS) -> Array: 1308 match r: 1309 case AS{a, +m}: 1310 lsm_r3(c, off, m, lsum(Nat.sub(c, 1n), off, m, 0.0, Array.get(F32, a, off))) 1311 1312 def lsm_row(+c: Nat, +off: U32, x: Array) -> Array: 1313 lsm_r2(c, off, pmax0(Nat.sub(c, 1n), off, Array.get(F32, x, off))) 1314 1315 def lsmf_rows(left: Nat, +i: U32, +c: Nat, x: Array) -> Array: 1316 match left: 1317 case 0n: 1318 lsm_row(c, (i * U32.from_nat(c) : U32), x) 1319 case 1n+q: 1320 lsmf_rows(q, (i + 1 : U32), c, lsm_row(c, (i * U32.from_nat(c) : U32), x)) 1321 1322 def lsm_fwd(+n: Nat, +c: Nat, x: TA.Mat) -> Fwd, TA.Mat>: 1323 match x: 1324 case TA.Mat{a}: 1325 out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{lsmf_rows(Nat.sub(n, 1n), 0, c, a)})) 1326 1327 # acc + a[j] for j = idx..idx + k, in order 1328 def asum(k: Nat, +idx: U32, acc: F32, r: Array & F32) -> AS: 1329 match k r: 1330 case 0n Tuple{a, +x}: 1331 AS{a, (acc + x : F32)} 1332 case 1n+q Tuple{a, +x}: 1333 asum(q, (idx + 1 : U32), (acc + x : F32), Array.get(F32, a, (idx + 1 : U32))) 1334 1335 # dy[j] = dy[j] - exp(y[j]) * s, j = idx..idx + k 1336 def lsmb_fix(k: Nat, +idx: U32, +s: F32, rd: Array & F32, ry: Array & F32) -> A2: 1337 match k rd ry: 1338 case 0n Tuple{d, +g} Tuple{y, +v}: 1339 A2{Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), y} 1340 case 1n+q Tuple{d, +g} Tuple{y, +v}: 1341 lsmb_fix(q, (idx + 1 : U32), s, Array.get(F32, Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), (idx + 1 : U32)), Array.get(F32, y, (idx + 1 : U32))) 1342 1343 def lsmb_row(+c: Nat, +off: U32, y: Array, r: AS) -> A2: 1344 match r: 1345 case AS{d, +s}: 1346 lsmb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off)) 1347 1348 # dx = dy - exp(y) * sum(dy), row after row, in place over dy 1349 def lsmb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2: 1350 match left st: 1351 case 0n A2{d, y}: 1352 lsmb_row(c, (i * U32.from_nat(c) : U32), y, asum(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), 0.0, Array.get(F32, d, (i * U32.from_nat(c) : U32)))) 1353 case 1n+q A2{d, y}: 1354 lsmb_rows(q, (i + 1 : U32), c, lsmb_row(c, (i * U32.from_nat(c) : U32), y, asum(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), 0.0, Array.get(F32, d, (i * U32.from_nat(c) : U32))))) 1355 1356 def lsm_bwd(+n: Nat, +c: Nat, s: TA.Mat, dy: TA.Mat) -> Bwd>: 1357 match s dy: 1358 case TA.Mat{y} TA.Mat{d}: 1359 dx_fin(n, c, lsmb_rows(Nat.sub(n, 1n), 0, c, A2{d, y})) 1360 1361 def LogSoftmax(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1362 Layer{p => x => lsm_fwd(n, c, x), p => s => dy => g => lsm_bwd(n, c, s, dy), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1363 1364 # ---- RMSNorm: y = g * (x / sqrt(mean(x^2) + eps)); its parameter is the gain g (1 x d) ---- 1365 1366 # one optimizer step on a layer whose parameter is a single matrix 1367 def mat_step_fin(-r: Nat, -c: Nat, q: MatStep) -> Step>: 1368 match q: 1369 case MatStep{w, g, st}: Step{w, g, st} 1370 1371 def mat_step(+r: Nat, +c: Nat, +op: Opt, p: TA.Mat, g: TA.Mat, st: OptState>) -> Step>: 1372 mat_step_fin(r, c, Opt.mat(r, c, op, p, g, st)) 1373 1374 def g_add_fin(-r: Nat, -c: Nat, r2: TA.MPair) -> TA.Mat: 1375 match r2: 1376 case TA.MPair{a, b}: 1377 a 1378 1379 # a + b, dropping b 1380 def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.Mat: 1381 g_add_fin(r, c, Mat.add(r, c, a, b)) 1382 1383 def dup_fin(-r: Nat, -c: Nat, m: TA.MM2) -> Two, TA.Mat>: 1384 match m: 1385 case TA.MM2{a, b}: 1386 Two{a, b} 1387 1388 # two copies of a matrix 1389 def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat) -> Two, TA.Mat>: 1390 dup_fin(r, c, TA.Mat.clone(r, c, m)) 1391 1392 def ml_list(-r: Nat, -c: Nat, ml: TA.ML) -> List<&2, F32>: 1393 match ml: 1394 case TA.ML{m, xs}: 1395 xs 1396 1397 # ---- RMSNorm layer, a row at a time over the arrays: x (n x d), the gain g (1 x d), the output y. Each sum is added 1398 # left to right from +0.0, each product in the order the list version multiplied, so the bits are the list version's. 1399 # ---- 1400 1401 # acc + a[i]·a[i] + ... over k + 1 numbers from i, in order; r holds a[i], already read 1402 def ssq(k: Nat, +i: U32, acc: F32, r: Array & F32) -> AS: 1403 match k r: 1404 case 0n Tuple{a, +x}: 1405 AS{a, (acc + (x * x : F32) : F32)} 1406 case 1n+p Tuple{a, +x}: 1407 ssq(p, (i + 1 : U32), (acc + (x * x : F32) : F32), Array.get(F32, a, (i + 1 : U32))) 1408 1409 # 1 / rms of a row whose sum of squares is s 1410 def rms_inv(+s: F32, +df: F32, +eps: F32) -> F32: 1411 (1.0 / F32.sqrt(((s / df : F32) + eps : F32)) : F32) 1412 1413 type R3 is Type: 1414 R3{x: Array, g: Array, y: Array} 1415 1416 # y[off + j] = g[j] * (x[off + j] * r), j = jj..jj + k; rx and rg hold x[off + jj] and g[jj] 1417 def rms_w(k: Nat, +j: U32, +off: U32, +r: F32, y: Array, rx: Array & F32, rg: Array & F32) -> R3: 1418 match k rx rg: 1419 case 0n Tuple{x, +a} Tuple{g, +b}: 1420 R3{x, g, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32))} 1421 case 1n+p Tuple{x, +a} Tuple{g, +b}: 1422 rms_w(p, (j + 1 : U32), off, r, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1423 1424 def rmsf_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, g: Array, y: Array, r: AS) -> R3: 1425 match r: 1426 case AS{x, +s}: 1427 rms_w(Nat.sub(d, 1n), 0, off, rms_inv(s, df, eps), y, Array.get(F32, x, off), Array.get(F32, g, 0)) 1428 1429 def rmsf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R3) -> R3: 1430 match s: 1431 case R3{x, g, y}: 1432 rmsf_row2(d, off, df, eps, g, y, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off))) 1433 1434 # rows i.. (left + 1 of them, d wide) 1435 def rmsf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R3) -> R3: 1436 match left: 1437 case 0n: 1438 rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1439 case 1n+p: 1440 rmsf_rows(p, (i + 1 : U32), d, df, eps, rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1441 1442 def rms_f2(-n: Nat, -d: Nat, s: R3) -> Fwd, TA.Mat, TA.Mat>: 1443 match s: 1444 case R3{x, g, y}: 1445 Fwd{TA.Mat{g}, TA.Mat{y}, TA.Mat{x}} 1446 1447 def rms_fwd(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 1448 match p x: 1449 case TA.Mat{g} TA.Mat{xa}: 1450 rms_f2(n, d, rmsf_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R3{xa, g, zeros_a(Nat.mul(n, d))})) 1451 1452 # with r = 1/rms, xh = x*r, dxh = dy*g: dx = r * (dxh - xh * mean(dxh * xh)); dg += dy * xh, row after row 1453 type R5 is Type: 1454 R5{x: Array, dy: Array, g: Array, dx: Array, dg: Array} 1455 1456 type R3S is Type: 1457 R3S{x: Array, dy: Array, g: Array, s: F32} 1458 1459 # acc + dxh[j]·xh[j] + ..., in order; rx, rd, rg hold x[off + jj], dy[off + jj], g[jj] 1460 def rmsb_dot(k: Nat, +j: U32, +off: U32, +r: F32, acc: F32, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R3S: 1461 match k rx rd rg: 1462 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1463 R3S{x, dy, g, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32)} 1464 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1465 rmsb_dot(p, (j + 1 : U32), off, r, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1466 1467 def dg_add(+j: U32, +v: F32, r: Array & F32) -> Array: 1468 match r: 1469 case Tuple{dg, +old}: 1470 Array.set(F32, dg, j, (old + v : F32)) 1471 1472 # dx[off + j] = ((dy[j]·g[j]) - (x[j]·r)·m)·r and dg[j] += dy[j]·(x[j]·r), j = jj..jj + k 1473 def rmsb_w(k: Nat, +j: U32, +off: U32, +r: F32, +m: F32, dx: Array, dg: Array, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R5: 1474 match k rx rd rg: 1475 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1476 +xh = (a * r : F32) 1477 R5{x, dy, g, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), dg_add(j, (b * xh : F32), Array.get(F32, dg, j))} 1478 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1479 +xh = (a * r : F32) 1480 rmsb_w(p, (j + 1 : U32), off, r, m, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), dg_add(j, (b * xh : F32), Array.get(F32, dg, j)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1481 1482 def rmsb_row3(+d: Nat, +off: U32, +df: F32, +r: F32, dx: Array, dg: Array, q: R3S) -> R5: 1483 match q: 1484 case R3S{x, dy, g, +s}: 1485 rmsb_w(Nat.sub(d, 1n), 0, off, r, (s / df : F32), dx, dg, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)) 1486 1487 def rmsb_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, dy: Array, g: Array, dx: Array, dg: Array, q: AS) -> R5: 1488 match q: 1489 case AS{x, +s}: 1490 +r = rms_inv(s, df, eps) 1491 rmsb_row3(d, off, df, r, dx, dg, rmsb_dot(Nat.sub(d, 1n), 0, off, r, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0))) 1492 1493 def rmsb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R5) -> R5: 1494 match s: 1495 case R5{x, dy, g, dx, dg}: 1496 rmsb_row2(d, off, df, eps, dy, g, dx, dg, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off))) 1497 1498 def rmsb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R5) -> R5: 1499 match left: 1500 case 0n: 1501 rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1502 case 1n+p: 1503 rmsb_rows(p, (i + 1 : U32), d, df, eps, rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1504 1505 def gain_bwd_fin(-n: Nat, -d: Nat, -P: Type, p: P, dx: TA.Mat, g: P) -> Bwd>: 1506 Bwd{p, dx, g} 1507 1508 def rms_b2(+d: Nat, -n: Nat, gacc: TA.Mat<1n, d>, s: R5) -> Bwd, TA.Mat>: 1509 match s: 1510 case R5{x, dy, g, dx, dg}: 1511 gain_bwd_fin(n, d, TA.Mat<1n, d>, TA.Mat{g}, TA.Mat{dx}, Mat.sum2(1n, d, gacc, TA.Mat{dg})) 1512 1513 def rms_bwd(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat, dy: TA.Mat, gacc: TA.Mat<1n, d>) -> Bwd, TA.Mat>: 1514 match p x dy: 1515 case TA.Mat{g} TA.Mat{xa} TA.Mat{da}: 1516 rms_b2(d, n, gacc, rmsb_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R5{xa, da, g, zeros_a(Nat.mul(n, d)), zeros_a(d)})) 1517 1518 # RMSNorm over rows of d numbers, with a learned gain (Llama's norm); saves its input 1519 def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 1520 Layer{p => x => rms_fwd(n, d, eps, p, x), p => s => dy => g => rms_bwd(n, d, eps, p, s, dy, g), op => p => g => st => mat_step(1n, d, op, p, g, st), u => TA.Mat.zeros(1n, d), p => Mat.dump(1n, d, p), p => xs => off => Mat.restore(1n, d, p, xs, off), p => xs => off => Mat.store(1n, d, p, xs, off), u => U32.from_nat(d)} 1521 1522 # ---- LayerNorm: y = (x - mean) / sqrt(var + eps) * g + b, var the biased variance ---- 1523 1524 type LN<-d: Nat> is Type: 1525 LN{g: TA.Mat<1n, d>, b: TA.Mat<1n, d>} 1526 1527 # acc + (a[j] - mu)·(a[j] - mu) for j = idx..idx + k, in order 1528 def csq(k: Nat, +idx: U32, +mu: F32, acc: F32, r: Array & F32) -> AS: 1529 match k r: 1530 case 0n Tuple{a, +x}: 1531 AS{a, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32)} 1532 case 1n+q Tuple{a, +x}: 1533 csq(q, (idx + 1 : U32), mu, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32), Array.get(F32, a, (idx + 1 : U32))) 1534 1535 type ASM is Type: 1536 ASM{a: Array, mu: F32, r: F32} 1537 1538 def ln_mr2(+df: F32, +eps: F32, +mu: F32, q: AS) -> ASM: 1539 match q: 1540 case AS{a, +v}: 1541 ASM{a, mu, (1.0 / F32.sqrt(((v / df : F32) + eps : F32)) : F32)} 1542 1543 def ln_mr1(+d: Nat, +off: U32, +df: F32, +eps: F32, q: AS) -> ASM: 1544 match q: 1545 case AS{a, +s}: 1546 ln_mr2(df, eps, (s / df : F32), csq(Nat.sub(d, 1n), off, (s / df : F32), 0.0, Array.get(F32, a, off))) 1547 1548 # a row's mean and 1 / sqrt(var + eps), as ln_xh and ln_r computed them 1549 def ln_mr(+d: Nat, +off: U32, +df: F32, +eps: F32, x: Array) -> ASM: 1550 ln_mr1(d, off, df, eps, asum(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off))) 1551 1552 type R4 is Type: 1553 R4{x: Array, g: Array, b: Array, y: Array} 1554 1555 # y[off + j] = ((x - mu) * r) * g[j] + b[j], j = jj..jj + k 1556 def lnf_w(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, y: Array, rx: Array & F32, rg: Array & F32, rb: Array & F32) -> R4: 1557 match k rx rg rb: 1558 case 0n Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}: 1559 R4{x, g, b, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32))} 1560 case 1n+q Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}: 1561 lnf_w(q, (j + 1 : U32), off, mu, r, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)), Array.get(F32, b, (j + 1 : U32))) 1562 1563 def lnf_row2(+d: Nat, +off: U32, g: Array, b: Array, y: Array, m: ASM) -> R4: 1564 match m: 1565 case ASM{x, +mu, +r}: 1566 lnf_w(Nat.sub(d, 1n), 0, off, mu, r, y, Array.get(F32, x, off), Array.get(F32, g, 0), Array.get(F32, b, 0)) 1567 1568 def lnf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4) -> R4: 1569 match s: 1570 case R4{x, g, b, y}: 1571 lnf_row2(d, off, g, b, y, ln_mr(d, off, df, eps, x)) 1572 1573 def lnf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4) -> R4: 1574 match left: 1575 case 0n: 1576 lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1577 case 1n+q: 1578 lnf_rows(q, (i + 1 : U32), d, df, eps, lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1579 1580 def ln_f2(-n: Nat, -d: Nat, s: R4) -> Fwd, TA.Mat, TA.Mat>: 1581 match s: 1582 case R4{x, g, b, y}: 1583 Fwd{LN{TA.Mat{g}, TA.Mat{b}}, TA.Mat{y}, TA.Mat{x}} 1584 1585 def ln_fwd(+n: Nat, +d: Nat, +eps: F32, p: LN, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 1586 match p x: 1587 case LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{xa}: 1588 ln_f2(n, d, lnf_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R4{xa, g, b, zeros_a(Nat.mul(n, d))})) 1589 1590 # with xh = (x - mu) * r and dxh = dy * g: dx = r * ((dxh - mean(dxh)) - xh * mean(dxh * xh)); dg += dy * xh; db += dy 1591 type R3SS is Type: 1592 R3SS{x: Array, dy: Array, g: Array, sa: F32, sb: F32} 1593 1594 # sa + dxh[j] and sb + dxh[j]·xh[j], j = jj..jj + k, each in order 1595 def lnb_sums(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, sa: F32, sb: F32, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R3SS: 1596 match k rx rd rg: 1597 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1598 R3SS{x, dy, g, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32)} 1599 case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1600 lnb_sums(q, (j + 1 : U32), off, mu, r, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1601 1602 def ln_dw(+j: U32, +v: F32, +w: F32, st: A2) -> A2: 1603 match st: 1604 case A2{dg, db}: 1605 A2{dg_add(j, v, Array.get(F32, dg, j)), dg_add(j, w, Array.get(F32, db, j))} 1606 1607 type R4B is Type: 1608 R4B{x: Array, dy: Array, g: Array, dx: Array, w: A2} 1609 1610 # dx[off + j] = (((dxh - ma) - xh * mb) * r), dg[j] += dy * xh, db[j] += dy, j = jj..jj + k (ma, mb the two means) 1611 def lnb_w(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, +ma: F32, +mb: F32, dx: Array, w: A2, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R4B: 1612 match k rx rd rg: 1613 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1614 +xh = ((a - mu : F32) * r : F32) 1615 R4B{x, dy, g, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), ln_dw(j, (b * xh : F32), b, w)} 1616 case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1617 +xh = ((a - mu : F32) * r : F32) 1618 lnb_w(q, (j + 1 : U32), off, mu, r, ma, mb, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), ln_dw(j, (b * xh : F32), b, w), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1619 1620 def lnb_row3(+d: Nat, +off: U32, +df: F32, +mu: F32, +r: F32, dx: Array, w: A2, q: R3SS) -> R4B: 1621 match q: 1622 case R3SS{x, dy, g, +sa, +sb}: 1623 lnb_w(Nat.sub(d, 1n), 0, off, mu, r, (sa / df : F32), (sb / df : F32), dx, w, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)) 1624 1625 def lnb_row2(+d: Nat, +off: U32, +df: F32, dy: Array, g: Array, dx: Array, w: A2, m: ASM) -> R4B: 1626 match m: 1627 case ASM{x, +mu, +r}: 1628 lnb_row3(d, off, df, mu, r, dx, w, lnb_sums(Nat.sub(d, 1n), 0, off, mu, r, 0.0, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0))) 1629 1630 def lnb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4B) -> R4B: 1631 match s: 1632 case R4B{x, dy, g, dx, w}: 1633 lnb_row2(d, off, df, dy, g, dx, w, ln_mr(d, off, df, eps, x)) 1634 1635 def lnb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4B) -> R4B: 1636 match left: 1637 case 0n: 1638 lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1639 case 1n+q: 1640 lnb_rows(q, (i + 1 : U32), d, df, eps, lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1641 1642 def ln_b2(+d: Nat, -n: Nat, b: TA.Mat<1n, d>, gg: LN, s: R4B) -> Bwd, TA.Mat>: 1643 match gg s: 1644 case LN{ga, gb} R4B{x, dy, g, dx, A2{dg, db}}: 1645 gain_bwd_fin(n, d, LN, LN{TA.Mat{g}, b}, TA.Mat{dx}, LN{Mat.sum2(1n, d, ga, TA.Mat{dg}), Mat.sum2(1n, d, gb, TA.Mat{db})}) 1646 1647 def ln_bwd(+n: Nat, +d: Nat, +eps: F32, p: LN, x: TA.Mat, dy: TA.Mat, gg: LN) -> Bwd, TA.Mat>: 1648 match p x dy: 1649 case LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da}: 1650 ln_b2(d, n, b, gg, lnb_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R4B{xa, da, g, zeros_a(Nat.mul(n, d)), A2{zeros_a(d), zeros_a(d)}})) 1651 1652 def ln_pair(-d: Nat, p: LN) -> Two, TA.Mat<1n, d>>: 1653 match p: 1654 case LN{w, b}: Two{w, b} 1655 1656 def ln_unpair(-d: Nat, p: Two, TA.Mat<1n, d>>) -> LN: 1657 match p: 1658 case Two{w, b}: LN{w, b} 1659 1660 def ln_s2(-d: Nat, a: MatStep<1n, d>, c: MatStep<1n, d>) -> Step>: 1661 match a c: 1662 case MatStep{w, gw, sw} MatStep{b, gb, sb}: 1663 Step{LN{w, b}, LN{gw, gb}, OptState.map(Two, TA.Mat<1n, d>>, LN, OptState.join(TA.Mat<1n, d>, TA.Mat<1n, d>, sw, sb), x => ln_unpair(d, x), x => ln_unpair(d, x))} 1664 1665 def ln_step2(+d: Nat, +op: Opt, p: LN, g: LN, st: Two>, OptState>>) -> Step>: 1666 match p g st: 1667 case LN{w, b} LN{gw, gb} Two{sw, sb}: 1668 ln_s2(d, Opt.mat(1n, d, op, w, gw, sw), Opt.mat(1n, d, op, b, gb, sb)) 1669 1670 def ln_step(+d: Nat, +op: Opt, p: LN, g: LN, st: OptState>) -> Step>: 1671 ln_step2(d, op, p, g, OptState.split(TA.Mat<1n, d>, TA.Mat<1n, d>, OptState.map(LN, Two, TA.Mat<1n, d>>, st, x => ln_pair(d, x), x => ln_pair(d, x)))) 1672 1673 def ln_dump2(-d: Nat, r: PL, TA.Mat<1n, d>>>) -> PL>: 1674 match r: 1675 case PL{Two{g, b}, xs}: 1676 PL{LN{g, b}, xs} 1677 1678 def ln_dump(+d: Nat, p: LN) -> PL>: 1679 match p: 1680 case LN{g, b}: 1681 ln_dump2(d, pl_pair(TA.Mat<1n, d>, TA.Mat<1n, d>, Mat.dump(1n, d, g), Mat.dump(1n, d, b))) 1682 1683 def ln_load3(-d: Nat, r: PA, TA.Mat<1n, d>>>) -> PA>: 1684 match r: 1685 case PA{Two{g, b}, xs, off}: PA{LN{g, b}, xs, off} 1686 1687 def ln_load2(+d: Nat, b: TA.Mat<1n, d>, r: PA>) -> PA>: 1688 match r: 1689 case PA{g, xs, off}: 1690 ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.restore(1n, d, b, xs, off))) 1691 1692 def ln_load(+d: Nat, p: LN, xs: Array, off: U32) -> PA>: 1693 match p: 1694 case LN{g, b}: ln_load2(d, b, Mat.restore(1n, d, g, xs, off)) 1695 1696 def ln_store2(+d: Nat, b: TA.Mat<1n, d>, r: PA>) -> PA>: 1697 match r: 1698 case PA{g, xs, off}: 1699 ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.store(1n, d, b, xs, off))) 1700 1701 def ln_store(+d: Nat, p: LN, xs: Array, off: U32) -> PA>: 1702 match p: 1703 case LN{g, b}: ln_store2(d, b, Mat.store(1n, d, g, xs, off)) 1704 1705 def LayerNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 1706 Layer{p => x => ln_fwd(n, d, eps, p, x), p => s => dy => g => ln_bwd(n, d, eps, p, s, dy, g), op => p => g => st => ln_step(d, op, p, g, st), u => LN{TA.Mat.zeros(1n, d), TA.Mat.zeros(1n, d)}, p => ln_dump(d, p), p => xs => off => ln_load(d, p, xs, off), p => xs => off => ln_store(d, p, xs, off), u => U32.from_nat(Nat.mul(2n, d))} 1707 1708 # ===================================================================================================== 1709 # Embedding: row i of the output is row ids[i] of the table (v x d); the gradient adds each row back 1710 # ===================================================================================================== 1711 1712 # row i of the output (d numbers) copied from row ids[i] of the table, for the first n ids 1713 def emb_cp(left: Nat, ids: List<&2, Nat>, +i: U32, +d: Nat, st: A2) -> A2: 1714 match left ids st: 1715 case 1n+q Con{+id, t} A2{e, o}: 1716 emb_cp(q, t, (i + 1 : U32), d, copy_n(Nat.sub(d, 1n), (U32.from_nat(id) * U32.from_nat(d) : U32), (i * U32.from_nat(d) : U32), o, Array.get(F32, e, (U32.from_nat(id) * U32.from_nat(d) : U32)))) 1717 case _ _ st: 1718 st 1719 1720 def emb_f2(-n: Nat, -v: Nat, -d: Nat, +ids: List<&2, Nat>, r: A2) -> Fwd, TA.Mat, List<&2, Nat>>: 1721 match r: 1722 case A2{e, o}: 1723 Fwd{TA.Mat{e}, TA.Mat{o}, ids} 1724 1725 def emb_f(+n: Nat, +v: Nat, +d: Nat, +ids: List<&2, Nat>, e: TA.Mat) -> Fwd, TA.Mat, List<&2, Nat>>: 1726 match e: 1727 case TA.Mat{ea}: 1728 emb_f2(n, v, d, ids, emb_cp(n, ids, 0, d, A2{ea, zeros_a(Nat.mul(n, d))})) 1729 1730 # g[gi + j] = g[gi + j] + dy[di + j], j = 0..k 1731 def emb_add(k: Nat, +gi: U32, +di: U32, rg: Array & F32, rd: Array & F32) -> A2: 1732 match k rg rd: 1733 case 0n Tuple{g, +a} Tuple{dy, +b}: 1734 A2{Array.set(F32, g, gi, (a + b : F32)), dy} 1735 case 1n+q Tuple{g, +a} Tuple{dy, +b}: 1736 emb_add(q, (gi + 1 : U32), (di + 1 : U32), Array.get(F32, Array.set(F32, g, gi, (a + b : F32)), (gi + 1 : U32)), Array.get(F32, dy, (di + 1 : U32))) 1737 1738 # each dY row added into the gradient's row ids[i], in order, by index 1739 def emb_sc(left: Nat, ids: List<&2, Nat>, +i: U32, +d: Nat, st: A2) -> A2: 1740 match left ids st: 1741 case 1n+q Con{+id, t} A2{g, dy}: 1742 emb_sc(q, t, (i + 1 : U32), d, emb_add(Nat.sub(d, 1n), (U32.from_nat(id) * U32.from_nat(d) : U32), (i * U32.from_nat(d) : U32), Array.get(F32, g, (U32.from_nat(id) * U32.from_nat(d) : U32)), Array.get(F32, dy, (i * U32.from_nat(d) : U32)))) 1743 case _ _ st: 1744 st 1745 1746 def emb_b2(-v: Nat, -d: Nat, e: TA.Mat, r: A2) -> Bwd, List<&2, Nat>>: 1747 match r: 1748 case A2{g, dy}: 1749 Bwd{e, Nil{}, TA.Mat{g}} 1750 1751 def emb_bwd(+n: Nat, +v: Nat, +d: Nat, e: TA.Mat, ids: List<&2, Nat>, dy: TA.Mat, g: TA.Mat) -> Bwd, List<&2, Nat>>: 1752 match dy g: 1753 case TA.Mat{da} TA.Mat{ga}: 1754 emb_b2(v, d, e, emb_sc(n, ids, 0, d, A2{ga, da})) 1755 1756 # a table of v rows of d numbers, looked up by n ids; an id has no gradient, so dX is the empty list 1757 def Embedding(+n: Nat, +v: Nat, +d: Nat) -> Layer, List<&2, Nat>, TA.Mat, List<&2, Nat>>: 1758 Layer{p => x => emb_f(n, v, d, x, p), p => s => dy => g => emb_bwd(n, v, d, p, s, dy, g), op => p => g => st => mat_step(v, d, op, p, g, st), u => TA.Mat.zeros(v, d), p => Mat.dump(v, d, p), p => xs => off => Mat.restore(v, d, p, xs, off), p => xs => off => Mat.store(v, d, p, xs, off), u => U32.from_nat(Nat.mul(v, d))} 1759 1760 # ===================================================================================================== 1761 # Linear without bias: y = x·Wᵀ (Llama's projections) 1762 # ===================================================================================================== 1763 1764 def lnb_f2(-n: Nat, -i: Nat, -o: Nat, r: TA.MMulNT) -> Fwd, TA.Mat, TA.Mat>: 1765 match r: 1766 case TA.MMulNT{x, w, y}: 1767 Fwd{w, y, x} 1768 1769 def lnb_b3(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat, dx: TA.Mat, r: TA.MMulTN) -> Bwd, TA.Mat>: 1770 match r: 1771 case TA.MMulTN{dy, x, g}: 1772 Bwd{w, dx, g} 1773 1774 def lnb_b2(+n: Nat, +i: Nat, +o: Nat, +par: Nat, x: TA.Mat, g: TA.Mat, r: TA.MMul) -> Bwd, TA.Mat>: 1775 match r: 1776 case TA.MMul{dy, w, dx}: 1777 lnb_b3(n, i, o, w, dx, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, g)) 1778 1779 def LinearNB(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 1780 Layer{w => x => lnb_f2(n, i, o, TA.Mat.matmul_nt(n, i, o, par, x, w)), w => s => dy => g => lnb_b2(n, i, o, par, s, g, TA.Mat.matmul_precise(n, o, i, par, dy, w)), op => p => g => st => mat_step(o, i, op, p, g, st), u => TA.Mat.zeros(o, i), p => Mat.dump(o, i, p), p => xs => off => Mat.restore(o, i, p, xs, off), p => xs => off => Mat.store(o, i, p, xs, off), u => U32.from_nat(Nat.mul(o, i))} 1781 1782 # ===================================================================================================== 1783 # Combinators: Id, Fork (one input to two layers), Both (two inputs, two layers), Add2, Mul2, Residual 1784 # ===================================================================================================== 1785 1786 def Id(~X: Type) -> Layer: 1787 Layer{p => x => Fwd{p, x, Unit{}}, p => s => dy => g => Bwd{p, dy, g}, op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1788 1789 # one optimizer step on two layers' parameters, of any types 1790 def pair_step2(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, +op: Opt, p: Two, g: Two, st: Two, OptState>) -> Step>: 1791 match p g st: 1792 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: 1793 pair_s3(P1, P2, Layer.step(P1, X1, Y1, S1, la, op, p1, g1, s1), Layer.step(P2, X2, Y2, S2, lb, op, p2, g2, s2)) 1794 1795 def pair_step(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, +op: Opt, p: Two, g: Two, st: OptState>) -> Step>: 1796 pair_step2(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, op, p, g, OptState.split(P1, P2, st)) 1797 1798 def pair_f4(-P1: Type, -P2: Type, -Y1: Type, -Y2: Type, -S1: Type, -S2: Type, p1: P1, y1: Y1, s1: S1, r: Fwd) -> Fwd, Two, Two>: 1799 match r: 1800 case Fwd{p2, y2, s2}: 1801 Fwd{Two{p1, p2}, Two{y1, y2}, Two{s1, s2}} 1802 1803 # runs lb on x2 once la's forward pass is done 1804 def pair_f3(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer, -P1: Type, -Y1: Type, -S1: Type, p2: P2, x2: X2, r: Fwd) -> Fwd, Two, Two>: 1805 match r: 1806 case Fwd{p1, y1, s1}: 1807 pair_f4(P1, P2, Y1, Y2, S1, S2, p1, y1, s1, Layer.fwd(P2, X2, Y2, S2, lb, p2, x2)) 1808 1809 # ---- Fork: x goes to both layers (copied by `dup`); the two input gradients are summed by `add` ---- 1810 1811 def fork_f2(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p1: P1, p2: P2, d: Two) -> Fwd, Two, Two>: 1812 match d: 1813 case Two{xa, xb}: 1814 pair_f3(~P2, ~X, ~Y2, ~S2, ~lb, P1, Y1, S1, p2, xb, Layer.fwd(P1, X, Y1, S1, la, p1, xa)) 1815 1816 def fork_b3(~X: Type, ~add: X -> X -> X, -P1: Type, -P2: Type, p2: P2, g2: P2, dx2: X, r: Bwd) -> Bwd, X>: 1817 match r: 1818 case Bwd{p1, dx1, g1}: 1819 Bwd{Two{p1, p2}, add(dx1, dx2), Two{g1, g2}} 1820 1821 def fork_b2(~P1: Type, ~X: Type, ~Y1: Type, ~S1: Type, ~add: X -> X -> X, ~la: Layer, -P2: Type, p1: P1, s1: S1, d1: Y1, g1: P1, r: Bwd) -> Bwd, X>: 1822 match r: 1823 case Bwd{p2, dx2, g2}: 1824 fork_b3(~X, ~add, P1, P2, p2, g2, dx2, Layer.bwd(P1, X, Y1, S1, la, p1, s1, d1, g1)) 1825 1826 def fork_bwd(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~add: X -> X -> X, ~la: Layer, ~lb: Layer, p: Two, s: Two, dy: Two, g: Two) -> Bwd, X>: 1827 match p s dy g: 1828 case Two{p1, p2} Two{s1, s2} Two{d1, d2} Two{g1, g2}: 1829 fork_b2(~P1, ~X, ~Y1, ~S1, ~add, ~la, P2, p1, s1, d1, g1, Layer.bwd(P2, X, Y2, S2, lb, p2, s2, d2, g2)) 1830 1831 def fork_fwd(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~dup: X -> Two, ~la: Layer, ~lb: Layer, p: Two, x: X) -> Fwd, Two, Two>: 1832 match p: 1833 case Two{p1, p2}: 1834 fork_f2(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p1, p2, dup(x)) 1835 1836 def Fork(~P1: Type, ~P2: Type, ~X: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~dup: X -> Two, ~add: X -> X -> X, ~la: Layer, ~lb: Layer) -> Layer, X, Two, Two>: 1837 Layer{p => x => fork_fwd(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~dup, ~la, ~lb, p, x), p => s => dy => g => fork_bwd(~P1, ~P2, ~X, ~Y1, ~Y2, ~S1, ~S2, ~add, ~la, ~lb, p, s, dy, g), op => p => g => st => pair_step(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P1, X, Y1, S1, la), Layer.zeros(P2, X, Y2, S2, lb)}, p => pair_dump(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X, ~Y1, ~S1, ~P2, ~X, ~Y2, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X, Y1, S1, la) + Layer.count(P2, X, Y2, S2, lb) : U32)} 1838 1839 # ---- Both: la on the first of a pair, lb on the second ---- 1840 1841 def both_b3(-P1: Type, -P2: Type, -X1: Type, -X2: Type, p2: P2, g2: P2, dx2: X2, r: Bwd) -> Bwd, Two>: 1842 match r: 1843 case Bwd{p1, dx1, g1}: 1844 Bwd{Two{p1, p2}, Two{dx1, dx2}, Two{g1, g2}} 1845 1846 def both_b2(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~la: Layer, -P2: Type, -X2: Type, p1: P1, s1: S1, d1: Y1, g1: P1, r: Bwd) -> Bwd, Two>: 1847 match r: 1848 case Bwd{p2, dx2, g2}: 1849 both_b3(P1, P2, X1, X2, p2, g2, dx2, Layer.bwd(P1, X1, Y1, S1, la, p1, s1, d1, g1)) 1850 1851 def both_bwd(~P1: Type, ~P2: Type, ~X1: Type, ~X2: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, s: Two, dy: Two, g: Two) -> Bwd, Two>: 1852 match p s dy g: 1853 case Two{p1, p2} Two{s1, s2} Two{d1, d2} Two{g1, g2}: 1854 both_b2(~P1, ~X1, ~Y1, ~S1, ~la, P2, X2, p1, s1, d1, g1, Layer.bwd(P2, X2, Y2, S2, lb, p2, s2, d2, g2)) 1855 1856 def both_fwd(~P1: Type, ~P2: Type, ~X1: Type, ~X2: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, x: Two) -> Fwd, Two, Two>: 1857 match p x: 1858 case Two{p1, p2} Two{x1, x2}: 1859 pair_f3(~P2, ~X2, ~Y2, ~S2, ~lb, P1, Y1, S1, p2, x2, Layer.fwd(P1, X1, Y1, S1, la, p1, x1)) 1860 1861 def Both(~P1: Type, ~P2: Type, ~X1: Type, ~X2: Type, ~Y1: Type, ~Y2: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer) -> Layer, Two, Two, Two>: 1862 Layer{p => x => both_fwd(~P1, ~P2, ~X1, ~X2, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p, x), p => s => dy => g => both_bwd(~P1, ~P2, ~X1, ~X2, ~Y1, ~Y2, ~S1, ~S2, ~la, ~lb, p, s, dy, g), op => p => g => st => pair_step(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P1, X1, Y1, S1, la), Layer.zeros(P2, X2, Y2, S2, lb)}, p => pair_dump(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X1, ~Y1, ~S1, ~P2, ~X2, ~Y2, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X1, Y1, S1, la) + Layer.count(P2, X2, Y2, S2, lb) : U32)} 1863 1864 # ---- Tie: first, mid, last, where first and last share one parameter (tied embeddings: the table is also the output 1865 # layer's weight). Forward runs the three in order with the same P for first and last; backward runs last, mid, first, 1866 # both adding into the same gradient, as PyTorch does for a tensor used twice. ---- 1867 1868 type Three<-A: Type, -B: Type, -C: Type> is Type: 1869 Three{a: A, b: B, c: C} 1870 1871 def tie_f4(-P: Type, -Q: Type, -W: Type, -S1: Type, -S2: Type, -S3: Type, q: Q, s1: S1, s2: S2, r: Fwd) -> Fwd, W, Three>: 1872 match r: 1873 case Fwd{p, w, s3}: 1874 Fwd{Two{p, q}, w, Three{s1, s2, s3}} 1875 1876 def tie_f3(~P: Type, ~Z: Type, ~W: Type, ~S3: Type, ~lc: Layer, -Q: Type, -S1: Type, -S2: Type, p: P, s1: S1, r: Fwd) -> Fwd, W, Three>: 1877 match r: 1878 case Fwd{q, z, s2}: 1879 tie_f4(P, Q, W, S1, S2, S3, q, s1, s2, Layer.fwd(P, Z, W, S3, lc, p, z)) 1880 1881 def tie_f2(~P: Type, ~Q: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S2: Type, ~S3: Type, ~lb: Layer, ~lc: Layer, -S1: Type, q: Q, r: Fwd) -> Fwd, W, Three>: 1882 match r: 1883 case Fwd{p, y, s1}: 1884 tie_f3(~P, ~Z, ~W, ~S3, ~lc, Q, S1, S2, p, s1, Layer.fwd(Q, Y, Z, S2, lb, q, y)) 1885 1886 def tie_fwd(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S1: Type, ~S2: Type, ~S3: Type, ~la: Layer, ~lb: Layer, ~lc: Layer, pq: Two, x: X) -> Fwd, W, Three>: 1887 match pq: 1888 case Two{p, q}: 1889 tie_f2(~P, ~Q, ~Y, ~Z, ~W, ~S2, ~S3, ~lb, ~lc, S1, q, Layer.fwd(P, X, Y, S1, la, p, x)) 1890 1891 def tie_b4(-P: Type, -Q: Type, -X: Type, q: Q, gq: Q, r: Bwd) -> Bwd, X>: 1892 match r: 1893 case Bwd{p, dx, gp}: 1894 Bwd{Two{p, q}, dx, Two{gp, gq}} 1895 1896 def tie_b3(~P: Type, ~X: Type, ~Y: Type, ~S1: Type, ~la: Layer, -Q: Type, p: P, s1: S1, gp: P, r: Bwd) -> Bwd, X>: 1897 match r: 1898 case Bwd{q, dy, gq}: 1899 tie_b4(P, Q, X, q, gq, Layer.bwd(P, X, Y, S1, la, p, s1, dy, gp)) 1900 1901 def tie_b2(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, q: Q, s1: S1, s2: S2, gq: Q, r: Bwd) -> Bwd, X>: 1902 match r: 1903 case Bwd{p, dz, gp}: 1904 tie_b3(~P, ~X, ~Y, ~S1, ~la, Q, p, s1, gp, Layer.bwd(Q, Y, Z, S2, lb, q, s2, dz, gq)) 1905 1906 def tie_bwd(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S1: Type, ~S2: Type, ~S3: Type, ~la: Layer, ~lb: Layer, ~lc: Layer, pq: Two, s: Three, dw: W, g: Two) -> Bwd, X>: 1907 match pq s g: 1908 case Two{p, q} Three{s1, s2, s3} Two{gp, gq}: 1909 tie_b2(~P, ~Q, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, q, s1, s2, gq, Layer.bwd(P, Z, W, S3, lc, p, s3, dw, gp)) 1910 1911 # The shared parameter steps once (through first's step; last's must be the same kind of layer for P, e.g. Embedding 1912 # and LinearNB over the same v x d matrix), the middle with its own. 1913 def Tie(~P: Type, ~Q: Type, ~X: Type, ~Y: Type, ~Z: Type, ~W: Type, ~S1: Type, ~S2: Type, ~S3: Type, ~la: Layer, ~lb: Layer, ~lc: Layer) -> Layer, X, W, Three>: 1914 Layer{pq => x => tie_fwd(~P, ~Q, ~X, ~Y, ~Z, ~W, ~S1, ~S2, ~S3, ~la, ~lb, ~lc, pq, x), pq => s => dw => g => tie_bwd(~P, ~Q, ~X, ~Y, ~Z, ~W, ~S1, ~S2, ~S3, ~la, ~lb, ~lc, pq, s, dw, g), op => p => g => st => pair_step(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, op, p, g, st), u => Two{Layer.zeros(P, X, Y, S1, la), Layer.zeros(Q, Y, Z, S2, lb)}, p => pair_dump(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P, ~X, ~Y, ~S1, ~Q, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P, X, Y, S1, la) + Layer.count(Q, Y, Z, S2, lb) : U32)} 1915 1916 # ---- Add2 and Mul2: element-wise sum and product of a pair of matrices ---- 1917 1918 def add2_f(+n: Nat, +c: Nat, x: Two, TA.Mat>) -> Fwd, Unit>: 1919 match x: 1920 case Two{a, b}: 1921 Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}} 1922 1923 def Add2(+n: Nat, +c: Nat) -> Layer, TA.Mat>, TA.Mat, Unit>: 1924 Layer{p => x => add2_f(n, c, x), p => s => dy => g => Bwd{p, Mat.dup(n, c, dy), g}, op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1925 1926 def mul2_f3(-n: Nat, -c: Nat, a2: TA.Mat, b2: TA.Mat, r: TA.MPair) -> Fwd, Two, TA.Mat>>: 1927 match r: 1928 case TA.MPair{y, b}: 1929 Fwd{Unit{}, y, Two{a2, b2}} 1930 1931 def mul2_f2(+n: Nat, +c: Nat, da: Two, TA.Mat>, db: Two, TA.Mat>) -> Fwd, Two, TA.Mat>>: 1932 match da db: 1933 case Two{a1, a2} Two{b1, b2}: 1934 mul2_f3(n, c, a2, b2, Ew.zip(~f32_mul, n, c, a1, b1)) 1935 1936 def mul2_f(+n: Nat, +c: Nat, x: Two, TA.Mat>) -> Fwd, Two, TA.Mat>>: 1937 match x: 1938 case Two{a, b}: 1939 mul2_f2(n, c, Mat.dup(n, c, a), Mat.dup(n, c, b)) 1940 1941 def mpair_a(-n: Nat, -c: Nat, r: TA.MPair) -> TA.Mat: 1942 match r: 1943 case TA.MPair{a, b}: 1944 a 1945 1946 # d(a*b) = (dy*b, dy*a) 1947 def mul2_b2(+n: Nat, +c: Nat, s: Two, TA.Mat>, d: Two, TA.Mat>) -> Bwd, TA.Mat>>: 1948 match s d: 1949 case Two{a, b} Two{d1, d2}: 1950 Bwd{Unit{}, Two{mpair_a(n, c, Ew.zip(~f32_mul, n, c, d1, b)), mpair_a(n, c, Ew.zip(~f32_mul, n, c, d2, a))}, Unit{}} 1951 1952 def Mul2(+n: Nat, +c: Nat) -> Layer, TA.Mat>, TA.Mat, Two, TA.Mat>>: 1953 Layer{p => x => mul2_f(n, c, x), p => s => dy => g => mul2_b2(n, c, s, Mat.dup(n, c, dy)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1954 1955 # ---- Residual: y = x + f(x), for any layer f from n x c to n x c ---- 1956 1957 def Residual(~n: Nat, ~c: Nat, ~P: Type, ~S: Type, ~f: Layer, TA.Mat, S>) -> Layer, Unit>, TA.Mat, TA.Mat, Two, Unit>>: 1958 Seq(~Two, ~Unit, ~TA.Mat, ~Two, TA.Mat>, ~TA.Mat, ~Two, ~Unit, ~Fork(~Unit, ~P, ~TA.Mat, ~TA.Mat, ~TA.Mat, ~Unit, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~Id(~TA.Mat), ~f), ~Add2(n, c)) 1959 1960 # ===================================================================================================== 1961 # RoPE (the rotate-half form of Llama and `transformers`): rows are positions (row r is position r mod t), each row 1962 # holds h heads of 2*half numbers, and the pair (x[c], x[c + half]) of a head is rotated by pos * inv[c], with 1963 # inv[c] = 1 / base^(2c / (2*half)). No parameters, nothing saved: the backward pass rotates by the opposite angle. 1964 # ===================================================================================================== 1965 1966 # Over the arrays: a table of cos and sin for each position mod t and each c (t x half), computed once a call with the 1967 # operations the list version used, then every pair rotated in place, so the bits are the list version's. 1968 1969 type CS is Type: 1970 CS{c: Array, s: Array} 1971 1972 # the angle of table entry e: inv[c] * position, position e / half, c = e mod half, inv[c] = 1 / base^(2c / (2*half)) 1973 def rope_ang(+e: U32, +half: U32, +base: F32) -> F32: 1974 ((1.0 / F32.pow(base, (U32.to_f32((U32.mod(e, half) + U32.mod(e, half) : U32)) / U32.to_f32((half + half : U32)) : F32)) : F32) * U32.to_f32(U32.div(e, half)) : F32) 1975 1976 # entries e.. of the tables, left more after e 1977 def rope_tab(left: Nat, +e: U32, +half: U32, +base: F32, tb: CS) -> CS: 1978 match left tb: 1979 case 0n CS{cs, sn}: 1980 +ang = rope_ang(e, half, base) 1981 CS{Array.set(F32, cs, e, F32.cos(ang)), Array.set(F32, sn, e, F32.sin(ang))} 1982 case 1n+q CS{cs, sn}: 1983 +ang = rope_ang(e, half, base) 1984 rope_tab(q, (e + 1 : U32), half, base, CS{Array.set(F32, cs, e, F32.cos(ang)), Array.set(F32, sn, e, F32.sin(ang))}) 1985 1986 type RS3 is Type: 1987 RS3{x: Array, c: Array, s: Array} 1988 1989 # one pair (a = x[i1], b = x[i2]) rotated in place; fwd picks the direction: (a cos - b sin, b cos + a sin) forward, 1990 # (a cos + b sin, b cos - a sin) backward, in PyTorch's order of operations for each 1991 def rope_p3(fwd: Bool, +i1: U32, +i2: U32, +a: F32, r2: Array & F32, rc: Array & F32, rs: Array & F32) -> RS3: 1992 match fwd r2 rc rs: 1993 case True{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}: 1994 RS3{Array.set(F32, Array.set(F32, x, i1, ((a * c : F32) + (F32.neg(b) * z : F32) : F32)), i2, ((b * c : F32) + (a * z : F32) : F32)), cs, sn} 1995 case False{} Tuple{x, +b} Tuple{cs, +c} Tuple{sn, +z}: 1996 RS3{Array.set(F32, Array.set(F32, x, i1, ((a * c : F32) + (b * z : F32) : F32)), i2, ((b * c : F32) + F32.neg((a * z : F32)) : F32)), cs, sn} 1997 1998 def rope_p2(+fwd: Bool, +i1: U32, +i2: U32, +k: U32, cs: Array, sn: Array, r1: Array & F32) -> RS3: 1999 match r1: 2000 case Tuple{x, +a}: Lines 1-2000 of 2994. More: https://jend.leeeet.dev/read.txt?path=bendtorch@1.0.1.0/main.bend&start_line=2001