main.bend
bend-ml-tensor-array@0.1.0.0 · 619 lines · 0 definitions
1# bend-ml-tensor-array: tensores sobre Array<F32> plano, com a shape no tipo.
2#
3# import bend-ml-tensor-array@0.1.0.0/main.bend as TA
4#
5# Mat<r, c> guarda r*c números num Array<F32> linha a linha (índice i*c + j). As
6# dimensões são parâmetros de tipo apagados: a mesma garantia do bend-ml-tensor
7# (produto com dimensões erradas não compila), mas ~50x mais rápido que listas:
8# Array.get/set por índice, e produtos em blocos de linhas em paralelo.
9#
10# Um Array é afim (um só dono): toda operação que lê um operando o devolve junto
11# com o resultado (como Array.get faz), em registros MM, AR, ... Use Mat.clone
12# quando precisar de duas cópias.
13#
14# O que NÃO é provado: a invariante "o Array tem capacidade >= r*c" vale porque
15# os construtores (Mat.zeros, Mat.of_list) alocam o tamanho certo; e F32 é
16# validado por testes contra o PyTorch (reference/test_tensor_array.py), não por prova.
17
18import Base
19
20# ---------------------------------------------------------------------
21# tipos de apoio
22# ---------------------------------------------------------------------
23
24type R is Type:
25 R{a: Array<F32>, b: Array<F32>, x: F32}
26
27type G is Type:
28 G{a: Array<F32>, b: Array<F32>, c: Array<F32>}
29
30# menor d com 2^d >= n (n >= 2)
31
32def dotg(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> R:
33 match k ra rb:
34 case 0n Tuple{a, x} Tuple{b, y}:
35 R{a, b, acc}
36 case 1n+p Tuple{a, +x} Tuple{b, +y}:
37 dotg(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)))
38
39# a linha i de C, colunas j.. ; r é o produto da coluna atual (já calculado)
40
41def gcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, +cidx: U32, c: Array<F32>, r: R) -> G:
42 match mleft r:
43 case 0n R{a, b, +x}:
44 G{a, b, Array.set(F32, c, cidx, x)}
45 case 1n+p R{a, b, +x}:
46 gcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, (cidx + 1 : U32), Array.set(F32, c, cidx, x), dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
47
48def grows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, +mu: U32, g: G) -> G:
49 match nleft g:
50 case 0n G{a, b, c}:
51 G{a, b, c}
52 case 1n+p G{a, b, c}:
53 grows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, mu, gcols(mm1, kn, (i * arow : U32), 0, sa, sb, bcol, (i * mu : U32), c, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0))))
54
55# C (n x m) = A · B con k passos; c é o array de saída (com capacidade n*m)
56
57def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
58 grows(n, Nat.sub(m, 1n), k, 0, arow, sa, sb, bcol, U32.from_nat(m), G{a, b, c})
59
60
61# ---------------------------------------------------------------------
62# gemm paralelo: 2^d blocos de linhas de C, cada um com a sua cópia de A e B,
63# devolvendo listas que são escritas em C no fim
64# ---------------------------------------------------------------------
65
66type S3 is Type:
67 S3{a: Array<F32>, b: Array<F32>, out: List<&2, F32>}
68
69def lcols(mleft: Nat, +kn: Nat, +ia0: U32, +ib0: U32, +sa: U32, +sb: U32, +bcol: U32, out: List<&2, F32>, r: R) -> S3:
70 match mleft r:
71 case 0n R{a, b, +x}:
72 S3{a, b, x <> out}
73 case 1n+p R{a, b, +x}:
74 lcols(p, kn, ia0, (ib0 + bcol : U32), sa, sb, bcol, x <> out, dotg(kn, ia0, (ib0 + bcol : U32), sa, sb, 0.0, Array.get(F32, a, ia0), Array.get(F32, b, (ib0 + bcol : U32))))
75
76def lrows(nleft: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, st: S3) -> S3:
77 match nleft st:
78 case 0n S3{a, b, out}:
79 S3{a, b, out}
80 case 1n+p S3{a, b, out}:
81 lrows(p, mm1, kn, (i + 1 : U32), arow, sa, sb, bcol, lcols(mm1, kn, (i * arow : U32), 0, sa, sb, bcol, out, dotg(kn, (i * arow : U32), 0, sa, sb, 0.0, Array.get(F32, a, (i * arow : U32)), Array.get(F32, b, 0))))
82
83def lblock_fin(st: S3) -> List<&2, F32>:
84 match st:
85 case S3{a, b, out}:
86 List.reverse(&2, F32, out)
87
88def lblock(cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>) -> List<&2, F32>:
89 lblock_fin(lrows(cnt, mm1, kn, i, arow, sa, sb, bcol, S3{a, b, Nil{}}))
90
91def lpar(d: Nat, +cnt: Nat, +mm1: Nat, +kn: Nat, +i: U32, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, ra: Array<F32> & Array<F32>, rb: Array<F32> & Array<F32>) -> List<&2, F32>:
92 match d ra rb:
93 case 0n Tuple{a, a2} Tuple{b, b2}:
94 lblock(cnt, mm1, kn, i, arow, sa, sb, bcol, a, b)
95 case 1n++q Tuple{a, a2} Tuple{b, b2}:
96 x y = lpar(q, Nat.div(cnt, 2n), mm1, kn, i, arow, sa, sb, bcol, Array.clone(F32, a), Array.clone(F32, b)) lpar(q, Nat.sub(cnt, Nat.div(cnt, 2n)), mm1, kn, (i + U32.from_nat(Nat.div(cnt, 2n)) : U32), arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2))
97 List.append(&2, F32, x, y)
98
99def write_l(xs: List<&2, F32>, +idx: U32, c: Array<F32>) -> Array<F32>:
100 match xs:
101 case Nil{}:
102 c
103 case Con{h, t}:
104 write_l(t, (idx + 1 : U32), Array.set(F32, c, idx, h))
105
106# c[idx..] += lista
107
108def gp3(a: Array<F32>, b: Array<F32>, c: Array<F32>, xs: List<&2, F32>) -> G:
109 G{a, b, write_l(xs, 0, c)}
110
111def gp2(d: Nat, +n: Nat, +mm1: Nat, +kn: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, c: Array<F32>, pa: Array<F32> & Array<F32>, pb: Array<F32> & Array<F32>) -> G:
112 match pa pb:
113 case Tuple{a1, a2} Tuple{b1, b2}:
114 gp3(a1, b1, c, lpar(d, n, mm1, kn, 0, arow, sa, sb, bcol, Array.clone(F32, a2), Array.clone(F32, b2)))
115
116# C (n x m) = A · B, com 2^d blocos de linhas em paralelo; devolve A, B (intactos) e C
117
118def gemm_par(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
119 gp2(d, n, Nat.sub(m, 1n), k, arow, sa, sb, bcol, c, Array.clone(F32, a), Array.clone(F32, b))
120
121# ---------------------------------------------------------------------
122# operações elementwise (em lugar, por índice)
123# ---------------------------------------------------------------------
124
125def depth(n: Nat) -> Nat:
126 Nat.add(U32.log2(U32.from_nat(Nat.sub(n, 1n))), 1n)
127
128# ---------------------------------------------------------------------
129# gemm: C[i][j] = soma_p A[i*arow + p*sa] * B[p*sb + j*bcol]
130# (nn: arow=K, sa=1, sb=M, bcol=1 | nt: B guardada m x k: sb=1, bcol=K | tn: A guardada k x n: arow=1, sa=N)
131# ---------------------------------------------------------------------
132
133# acc += A[ia + t*sa] * B[ib + t*sb] por k passos; ra, rb já trazem o primeiro par
134
135type P2 is Type:
136 P2{a: Array<F32>, b: Array<F32>}
137
138type AL is Type:
139 AL{c: Array<F32>, l: List<&2, F32>}
140
141# lista de m elementos a partir de idx (lê com Array.get, devolvendo o array)
142
143def read_l(mleft: Nat, +idx: U32, acc: List<&2, F32>, rc: Array<F32> & F32) -> AL:
144 match mleft rc:
145 case 0n Tuple{c, +v}:
146 AL{c, List.reverse(&2, F32, v <> acc)}
147 case 1n+p Tuple{c, +v}:
148 read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, c, (idx + 1 : U32)))
149
150# escreve a lista a partir de idx
151
152def addl(xs: List<&2, F32>, +idx: U32, rc: Array<F32> & F32) -> Array<F32>:
153 match xs rc:
154 case Nil{} Tuple{c, v}:
155 c
156 case Con{+h, t} Tuple{c, +v}:
157 addl(t, (idx + 1 : U32), Array.get(F32, Array.set(F32, c, idx, (v + h : F32)), (idx + 1 : U32)))
158
159# soma bias (lista de m) a cada uma das n linhas de c (n x m)
160
161def bias_rows(nleft: Nat, +i: U32, +mu: U32, +bl: List<&2, F32>, c: Array<F32>) -> Array<F32>:
162 match nleft:
163 case 0n:
164 c
165 case 1n+p:
166 bias_rows(p, (i + 1 : U32), mu, bl, addl(bl, (i * mu : U32), Array.get(F32, c, (i * mu : U32))))
167
168# relu em lugar nos primeiros n elementos
169
170def relu_l(nleft: Nat, +t: U32, rc: Array<F32> & F32) -> Array<F32>:
171 match nleft rc:
172 case 0n Tuple{c, v}:
173 c
174 case 1n+p Tuple{c, +v}:
175 relu_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, c, t, F32.max(v, 0.0)), (t + 1 : U32)))
176
177def pos2(b: Bool) -> F32:
178 match b:
179 case True{}:
180 1.0
181 case False{}:
182 0.0
183
184def pos(+y: F32) -> F32:
185 pos2(F32.is_gt(y, 0.0))
186
187# d[t] *= (h[t] > 0): gradiente da relu; devolve os dois arrays
188
189def mask_l(nleft: Nat, +t: U32, rd: Array<F32> & F32, rh: Array<F32> & F32) -> P2:
190 match nleft rd rh:
191 case 0n Tuple{d, x} Tuple{h, y}:
192 P2{d, h}
193 case 1n+p Tuple{d, +x} Tuple{h, +y}:
194 mask_l(p, (t + 1 : U32), Array.get(F32, Array.set(F32, d, t, (x * pos(y) : F32)), (t + 1 : U32)), Array.get(F32, h, (t + 1 : U32)))
195
196# w[t] -= lr * dw[t]
197
198def maxl2(xs: List<&2, F32>, cur: F32) -> F32:
199 match xs:
200 case Nil{}:
201 cur
202 case Con{+h, t}:
203 maxl2(t, F32.max(cur, h))
204
205def maxl(xs: List<&2, F32>) -> F32:
206 match xs:
207 case Nil{}:
208 0.0
209 case Con{+h, t}:
210 maxl2(t, h)
211
212def expl(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
213 match xs:
214 case Nil{}:
215 Nil{}
216 case Con{+h, t}:
217 F32.exp((h - m : F32)) <> expl(t, m)
218
219def suml(xs: List<&2, F32>, acc: F32) -> F32:
220 match xs:
221 case Nil{}:
222 acc
223 case Con{+h, t}:
224 suml(t, (acc + h : F32))
225
226def divl(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
227 match xs:
228 case Nil{}:
229 Nil{}
230 case Con{+h, t}:
231 (h / s : F32) <> divl(t, s)
232
233def softmax2(+es: List<&2, F32>) -> List<&2, F32>:
234 divl(es, suml(es, 0.0))
235
236def softmax(+xs: List<&2, F32>) -> List<&2, F32>:
237 softmax2(expl(xs, maxl(xs)))
238
239# p[label]
240
241def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array<F32> & F32, rd: Array<F32> & F32) -> P2:
242 match nleft rw rd:
243 case 0n Tuple{w, x} Tuple{d, y}:
244 P2{w, d}
245 case 1n+p Tuple{w, +x} Tuple{d, +y}:
246 sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x - (lr * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, d, (t + 1 : U32)))
247
248# ---------------------------------------------------------------------
249# softmax + entropia cruzada por linha (listas de 10 elementos: custo desprezível)
250# ---------------------------------------------------------------------
251
252def pick(ps: List<&2, F32>, n: Nat) -> F32:
253 match ps n:
254 case Con{+h, t} 0n:
255 h
256 case Con{h, t} 1n+q:
257 pick(t, q)
258 case Nil{} _:
259 0.0
260
261# (p - onehot(label)) / n
262
263def hit(b: Bool) -> F32:
264 match b:
265 case True{}:
266 1.0
267 case False{}:
268 0.0
269
270def grad_l(ps: List<&2, F32>, +label: Nat, +j: Nat, +nf: F32) -> List<&2, F32>:
271 match ps:
272 case Nil{}:
273 Nil{}
274 case Con{+h, t}:
275 (((h - hit(Nat.is_eq(label, j)) : F32) / nf) : F32) <> grad_l(t, label, 1n+j, nf)
276
277type L2 is Type:
278 L2{loss: F32, c: Array<F32>}
279
280def ce_row2(+ps: List<&2, F32>, +label: Nat, +nf: F32, +idx: U32, loss: F32, c: Array<F32>) -> L2:
281 L2{(loss - F32.log(F32.max(pick(ps, label), 0.000000000001)) : F32), write_l(grad_l(ps, label, 0n, nf), idx, c)}
282
283def ce_row(al: AL, +label: Nat, +nf: F32, +idx: U32, loss: F32) -> L2:
284 match al:
285 case AL{c, l}:
286 ce_row2(softmax(l), label, nf, idx, loss, c)
287
288# percorre as linhas de z2 (n x m): r traz a perda acumulada e o array; ao fim, o array virou o gradiente
289
290type Mat<-r: Nat, -c: Nat> is Type:
291 Mat{d: Array<F32>}
292
293# menor d com 2^d >= n (n >= 2); n < 2 vale 0
294
295def ce_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +nf: F32, +i: U32, r: L2) -> L2:
296 match labels r:
297 case Nil{} L2{loss, c}:
298 L2{loss, c}
299 case Con{+y, t} L2{loss, c}:
300 ce_rows(t, mm1, mu, nf, (i + 1 : U32), ce_row(read_l(mm1, (i * mu : U32), Nil{}, Array.get(F32, c, (i * mu : U32))), y, nf, (i * mu : U32), loss))
301
302# ---- apoio: acertos por linha ----
303
304type CR is Type:
305 CR{c: Array<F32>, n: Nat}
306
307def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
308 match xs:
309 case Nil{}:
310 False{}
311 case Con{+h, t}:
312 F32.is_gt(h, v)
313
314def argmax2(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> Nat:
315 match xs g:
316 case Nil{} _:
317 bi
318 case Con{+h, +t} True{}:
319 argmax2(t, gt_head(t, h), 1n+i, i, h)
320 case Con{h, +t} False{}:
321 argmax2(t, gt_head(t, bv), 1n+i, bi, bv)
322
323def argmax(xs: List<&2, F32>) -> Nat:
324 match xs:
325 case Nil{}:
326 0n
327 case Con{+h, +t}:
328 argmax2(t, gt_head(t, h), 1n, 0n, h)
329
330def one(b: Bool) -> Nat:
331 match b:
332 case True{}:
333 1n
334 case False{}:
335 0n
336
337def count_rows3(+y: Nat, n: Nat, al: AL) -> CR:
338 match al:
339 case AL{c, l}:
340 CR{c, Nat.add(n, one(Nat.is_eq(argmax(l), y)))}
341
342def count_rows(labels: List<&2, Nat>, +mm1: Nat, +mu: U32, +i: U32, r: CR) -> CR:
343 match labels r:
344 case Nil{} CR{c, n}:
345 CR{c, n}
346 case Con{+y, t} CR{c, n}:
347 count_rows(t, mm1, mu, (i + 1 : U32), count_rows3(y, n, read_l(mm1, (i * mu : U32), Nil{}, Array.get(F32, c, (i * mu : U32)))))
348
349
350# =====================================================================
351# API tipada: Mat<r, c> (as dimensões entram no tipo)
352# =====================================================================
353
354def cap_depth2(small: Bool, n: Nat) -> Nat:
355 match small:
356 case True{}:
357 0n
358 case False{}:
359 depth(n)
360
361def cap_depth(+n: Nat) -> Nat:
362 cap_depth2(Nat.is_lt(n, 2n), n)
363
364# r x c de zeros, com a capacidade exata (potência de 2 mínima)
365
366def Mat.zeros(+r: Nat, +c: Nat) -> Mat<r, c>:
367 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), 0.0)}
368
369def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat<r, c>:
370 Mat{Array.new(F32, cap_depth(Nat.mul(r, c)), v)}
371
372def fill_list(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
373 match xs:
374 case Nil{}:
375 a
376 case Con{h, t}:
377 fill_list(t, (i + 1 : U32), Array.set(F32, a, i, h))
378
379def len_is(xs: List<&2, F32>, n: Nat) -> Bool:
380 match xs n:
381 case Nil{} 0n:
382 True{}
383 case Con{h, t} 1n+p:
384 len_is(t, p)
385 case _ _:
386 False{}
387
388def Mat.of_list2(-r: Nat, -c: Nat, +rr: Nat, +cc: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
389 match ok:
390 case True{}:
391 Some{Mat{fill_list(xs, 0, Array.new(F32, cap_depth(Nat.mul(rr, cc)), 0.0))}}
392 case False{}:
393 None{}
394
395# Mat<r, c> a partir de r*c números em ordem de linha; None se o tamanho não bate
396
397def Mat.of_list(+r: Nat, +c: Nat, +xs: List<&2, F32>) -> Maybe<&1, Mat<r, c>>:
398 Mat.of_list2(r, c, r, c, len_is(xs, Nat.mul(r, c)), xs)
399
400# lê os n primeiros números, devolvendo a matriz
401
402type ML<-r: Nat, -c: Nat> is Type:
403 ML{m: Mat<r, c>, l: List<&2, F32>}
404
405def Mat.to_list2(-r: Nat, -c: Nat, al: AL) -> ML<r, c>:
406 match al:
407 case AL{d, l}:
408 ML{Mat{d}, l}
409
410# duas cópias independentes
411
412def Mat.to_list(+r: Nat, +c: Nat, m: Mat<r, c>) -> ML<r, c>:
413 match m:
414 case Mat{d}:
415 Mat.to_list2(r, c, read_l(Nat.sub(Nat.mul(r, c), 1n), 0, Nil{}, Array.get(F32, d, 0)))
416
417type MM2<-r: Nat, -c: Nat> is Type:
418 MM2{a: Mat<r, c>, b: Mat<r, c>}
419
420def Mat.clone2(-r: Nat, -c: Nat, p: Array<F32> & Array<F32>) -> MM2<r, c>:
421 match p:
422 case Tuple{x, y}:
423 MM2{Mat{x}, Mat{y}}
424
425# ---- produtos: C (n x m) = A · B, com 2^par blocos de linhas em paralelo ----
426
427# C = A(n x k) · B(k x m)
428
429def Mat.clone(-r: Nat, -c: Nat, m: Mat<r, c>) -> MM2<r, c>:
430 match m:
431 case Mat{d}:
432 Mat.clone2(r, c, Array.clone(F32, d))
433
434type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
435 MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}
436
437# C = A(n x k) · Bᵀ, com B guardada m x k
438
439type MMulNT<-n: Nat, -k: Nat, -m: Nat> is Type:
440 MMulNT{a: Mat<n, k>, b: Mat<m, k>, c: Mat<n, m>}
441
442# C = Aᵀ · B, com A guardada k x n
443
444type MMulTN<-n: Nat, -k: Nat, -m: Nat> is Type:
445 MMulTN{a: Mat<k, n>, b: Mat<k, m>, c: Mat<n, m>}
446
447def gemm_any(d: Nat, +n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32, a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:
448 match d:
449 case 0n:
450 gemm(n, m, k, arow, sa, sb, bcol, a, b, c)
451 case 1n+q:
452 gemm_par(1n+q, n, m, k, arow, sa, sb, bcol, a, b, c)
453
454def mmul_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMul<n, k, m>:
455 match g:
456 case G{a, b, c}:
457 MMul{Mat{a}, Mat{b}, Mat{c}}
458
459def mmul_nt_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulNT<n, k, m>:
460 match g:
461 case G{a, b, c}:
462 MMulNT{Mat{a}, Mat{b}, Mat{c}}
463
464def mmul_tn_fin(-n: Nat, -k: Nat, -m: Nat, g: G) -> MMulTN<n, k, m>:
465 match g:
466 case G{a, b, c}:
467 MMulTN{Mat{a}, Mat{b}, Mat{c}}
468
469def mm_a(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMul<n, k, m>:
470 mmul_fin(n, k, m, gemm_any(d, n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
471
472# C = A · B (n x k) · (k x m); par = log2 do número de blocos paralelos (0 = sequencial)
473
474def Mat.matmul(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<k, m>) -> MMul<n, k, m>:
475 match a b:
476 case Mat{x} Mat{y}:
477 mm_a(n, k, m, par, x, y)
478
479def mm_nt(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulNT<n, k, m>:
480 mmul_nt_fin(n, k, m, gemm_any(d, n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
481
482# C = A · Bᵀ (n x k) · (m x k)ᵀ: para pesos guardados com uma linha por saída
483
484def Mat.matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<n, k>, b: Mat<m, k>) -> MMulNT<n, k, m>:
485 match a b:
486 case Mat{x} Mat{y}:
487 mm_nt(n, k, m, par, x, y)
488
489def mm_tn(+n: Nat, +k: Nat, +m: Nat, d: Nat, a: Array<F32>, b: Array<F32>) -> MMulTN<n, k, m>:
490 mmul_tn_fin(n, k, m, gemm_any(d, n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, a, b, Array.new(F32, cap_depth(Nat.mul(n, m)), 0.0)))
491
492# C = Aᵀ · B (k x n)ᵀ · (k x m): gradiente de pesos, X^T · dY
493
494def Mat.matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: Mat<k, n>, b: Mat<k, m>) -> MMulTN<n, k, m>:
495 match a b:
496 case Mat{x} Mat{y}:
497 mm_tn(n, k, m, par, x, y)
498
499# ---- operações elementwise e do treino ----
500
501# soma o bias (uma linha, m números) a cada linha de C (n x m)
502
503type MBias<-n: Nat, -m: Nat> is Type:
504 MBias{c: Mat<n, m>, b: Mat<1n, m>}
505
506def add_row2(-n: Nat, -m: Nat, +nn: Nat, +mm: Nat, cd: Array<F32>, al: AL) -> MBias<n, m>:
507 match al:
508 case AL{bd, bl}:
509 MBias{Mat{bias_rows(nn, 0, U32.from_nat(mm), bl, cd)}, Mat{bd}}
510
511def Mat.add_row(+n: Nat, +m: Nat, c: Mat<n, m>, b: Mat<1n, m>) -> MBias<n, m>:
512 match c b:
513 case Mat{cd} Mat{bd}:
514 add_row2(n, m, n, m, cd, read_l(Nat.sub(m, 1n), 0, Nil{}, Array.get(F32, bd, 0)))
515
516# par de matrizes do mesmo formato (o resultado de uma operação que lê dois operandos)
517
518type MPair<-r: Nat, -c: Nat> is Type:
519 MPair{a: Mat<r, c>, b: Mat<r, c>}
520
521def mpair_fin(-r: Nat, -c: Nat, p: P2) -> MPair<r, c>:
522 match p:
523 case P2{a, b}:
524 MPair{Mat{a}, Mat{b}}
525
526def relu_fin(-r: Nat, -c: Nat, d: Array<F32>) -> Mat<r, c>:
527 Mat{d}
528
529def Mat.relu(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<r, c>:
530 match a:
531 case Mat{d}:
532 relu_fin(r, c, relu_l(Nat.mul(r, c), 0, Array.get(F32, d, 0)))
533
534# (dy * (h > 0), h): gradiente que atravessa a relu
535
536def Mat.relu_bwd(+r: Nat, +c: Nat, dy: Mat<r, c>, h: Mat<r, c>) -> MPair<r, c>:
537 match dy h:
538 case Mat{d} Mat{x}:
539 mpair_fin(r, c, mask_l(Nat.mul(r, c), 0, Array.get(F32, d, 0), Array.get(F32, x, 0)))
540
541# (w - lr * dw, dw): um passo de SGD
542
543def Mat.sgd(+r: Nat, +c: Nat, +lr: F32, w: Mat<r, c>, dw: Mat<r, c>) -> MPair<r, c>:
544 match w dw:
545 case Mat{x} Mat{y}:
546 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, x, 0), Array.get(F32, y, 0)))
547
548# (a + b, b): soma elementwise (a - (-1)*b é exatamente a + b)
549
550def Mat.add(+r: Nat, +c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> MPair<r, c>:
551 match a b:
552 case Mat{x} Mat{y}:
553 mpair_fin(r, c, sgd_l(Nat.mul(r, c), 0, (0.0 - 1.0 : F32), Array.get(F32, x, 0), Array.get(F32, y, 0)))
554
555# soma de cada coluna (n x m -> 1 x m), devolvendo a matriz
556
557type MColSum<-n: Nat, -m: Nat> is Type:
558 MColSum{a: Mat<n, m>, s: Mat<1n, m>}
559
560def colsum_fin(-n: Nat, -m: Nat, g: G) -> MColSum<n, m>:
561 match g:
562 case G{ones, a, s}:
563 MColSum{Mat{a}, Mat{s}}
564
565def Mat.col_sums(+n: Nat, +m: Nat, a: Mat<n, m>) -> MColSum<n, m>:
566 match a:
567 case Mat{d}:
568 colsum_fin(n, m, gemm(1n, m, n, 0, 1, U32.from_nat(m), 1, Array.new(F32, cap_depth(n), 1.0), d, Array.new(F32, cap_depth(m), 0.0)))
569
570# entropia cruzada média e seu gradiente nos logits: (softmax - one-hot) / n.
571# labels tem uma entrada por linha (n); o gradiente substitui os logits.
572
573type MCE<-n: Nat, -c: Nat> is Type:
574 MCE{loss: F32, grad: Mat<n, c>}
575
576def ce_fin(-n: Nat, -c: Nat, +nf: F32, r: L2) -> MCE<n, c>:
577 match r:
578 case L2{loss, d}:
579 MCE{(loss / nf : F32), Mat{d}}
580
581def Mat.softmax_ce(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MCE<n, c>:
582 match z:
583 case Mat{d}:
584 ce_fin(n, c, F32.from_nat(n), ce_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), F32.from_nat(n), 0, L2{0.0, d}))
585
586# acertos de argmax por linha contra labels
587
588type MHits<-n: Nat, -c: Nat> is Type:
589 MHits{z: Mat<n, c>, hits: Nat}
590
591def hits_fin(-n: Nat, -c: Nat, r: CR) -> MHits<n, c>:
592 match r:
593 case CR{d, k}:
594 MHits{Mat{d}, k}
595
596def Mat.count_correct(+n: Nat, +c: Nat, z: Mat<n, c>, labels: List<&2, Nat>) -> MHits<n, c>:
597 match z:
598 case Mat{d}:
599 hits_fin(n, c, count_rows(labels, Nat.sub(c, 1n), U32.from_nat(c), 0, CR{d, 0n}))
600
601# uma linha como lista (devolvendo a matriz) e o contrário
602
603type MRow<-r: Nat, -c: Nat> is Type:
604 MRow{m: Mat<r, c>, row: List<&2, F32>}
605
606def row_fin(-r: Nat, -c: Nat, al: AL) -> MRow<r, c>:
607 match al:
608 case AL{d, l}:
609 MRow{Mat{d}, l}
610
611def Mat.read_row(+r: Nat, +c: Nat, m: Mat<r, c>, +i: U32) -> MRow<r, c>:
612 match m:
613 case Mat{d}:
614 row_fin(r, c, read_l(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), Nil{}, Array.get(F32, d, (i * U32.from_nat(c) : U32))))
615
616def Mat.write_row(-r: Nat, +c: Nat, m: Mat<r, c>, +i: U32, xs: List<&2, F32>) -> Mat<r, c>:
617 match m:
618 case Mat{d}:
619 Mat{write_l(xs, (i * U32.from_nat(c) : U32), d)}