main.bend
bend-ml-tensor@0.1.1.0 · 505 lines · 0 definitions
1# bend-ml-tensor: tensores com a shape no TIPO. Erro de shape é erro de tipo.
2#
3# import bend-ml-tensor@0.1.0.0/main.bend as T
4#
5# Vec<n> é um vetor de n números F32; Mat<r, c> é uma matriz r x c (lista de r
6# linhas com c números). As dimensões são parâmetros de tipo apagados: não custam
7# nada em tempo de execução, mas o checker as confere em toda operação.
8
9import Base
10import bend-ml-nat-lemmas@0.1.0.0/main.bend as NL
11
12type Vec<-n: Nat> is Data:
13 Vec{xs: List<&2, F32>}
14
15type Mat<-r: Nat, -c: Nat> is Data:
16 Mat{rows: List<&2, List<&2, F32>>}
17
18# ---------------------------------------------------------------------
19# listas de F32 (internas; as operações públicas fixam as dimensões)
20# ---------------------------------------------------------------------
21
22def zip_with.add(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>:
23 match xs ys:
24 case Con{+h, t} Con{+g, u}:
25 (h + g : F32) <> zip_with.add(t, u)
26 case _ _:
27 Nil{}
28
29def zip_with.sub(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>:
30 match xs ys:
31 case Con{+h, t} Con{+g, u}:
32 (h - g : F32) <> zip_with.sub(t, u)
33 case _ _:
34 Nil{}
35
36def zip_with.mul(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>:
37 match xs ys:
38 case Con{+h, t} Con{+g, u}:
39 (h * g : F32) <> zip_with.mul(t, u)
40 case _ _:
41 Nil{}
42
43def scale.l(+k: F32, xs: List<&2, F32>) -> List<&2, F32>:
44 match xs:
45 case Nil{}:
46 Nil{}
47 case Con{+h, t}:
48 (k * h : F32) <> scale.l(k, t)
49
50def dot.go(xs: List<&2, F32>, ys: List<&2, F32>, acc: F32) -> F32:
51 match xs ys:
52 case Con{+h, t} Con{+g, u}:
53 dot.go(t, u, (acc + (h * g : F32) : F32))
54 case _ _:
55 acc
56
57def sum.go(xs: List<&2, F32>, acc: F32) -> F32:
58 match xs:
59 case Nil{}:
60 acc
61 case Con{+h, t}:
62 sum.go(t, (acc + h : F32))
63
64def relu.l(xs: List<&2, F32>) -> List<&2, F32>:
65 match xs:
66 case Nil{}:
67 Nil{}
68 case Con{+h, t}:
69 F32.max(h, 0.0) <> relu.l(t)
70
71def fill.l(n: Nat, +v: F32) -> List<&2, F32>:
72 match n:
73 case 0n:
74 Nil{}
75 case 1n+p:
76 v <> fill.l(p, v)
77
78# ---------------------------------------------------------------------
79# Vec
80# ---------------------------------------------------------------------
81
82def Vec.zeros(n: Nat) -> Vec<n>:
83 Vec{fill.l(n, 0.0)}
84
85def Vec.fill(n: Nat, +v: F32) -> Vec<n>:
86 Vec{fill.l(n, v)}
87
88def Vec.add(-n: Nat, a: Vec<n>, b: Vec<n>) -> Vec<n>:
89 match a b:
90 case Vec{xs} Vec{ys}:
91 Vec{zip_with.add(xs, ys)}
92
93def Vec.sub(-n: Nat, a: Vec<n>, b: Vec<n>) -> Vec<n>:
94 match a b:
95 case Vec{xs} Vec{ys}:
96 Vec{zip_with.sub(xs, ys)}
97
98def Vec.mul(-n: Nat, a: Vec<n>, b: Vec<n>) -> Vec<n>:
99 match a b:
100 case Vec{xs} Vec{ys}:
101 Vec{zip_with.mul(xs, ys)}
102
103def Vec.scale(-n: Nat, +k: F32, a: Vec<n>) -> Vec<n>:
104 match a:
105 case Vec{xs}:
106 Vec{scale.l(k, xs)}
107
108def Vec.dot(-n: Nat, a: Vec<n>, b: Vec<n>) -> F32:
109 match a b:
110 case Vec{xs} Vec{ys}:
111 dot.go(xs, ys, 0.0)
112
113def Vec.sum(-n: Nat, a: Vec<n>) -> F32:
114 match a:
115 case Vec{xs}:
116 sum.go(xs, 0.0)
117
118def Vec.relu(-n: Nat, a: Vec<n>) -> Vec<n>:
119 match a:
120 case Vec{xs}:
121 Vec{relu.l(xs)}
122
123# ---------------------------------------------------------------------
124# listas de linhas (internas)
125# ---------------------------------------------------------------------
126
127def heads(rows: List<&2, List<&2, F32>>) -> List<&2, F32>:
128 match rows:
129 case Con{Con{+h, t}, rest}:
130 h <> heads(rest)
131 case _:
132 Nil{}
133
134def tails(rows: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
135 match rows:
136 case Con{Con{h, t}, rest}:
137 t <> tails(rest)
138 case _:
139 Nil{}
140
141# transposta de uma matriz com `c` colunas (c é quem encolhe, vem primeiro)
142def transpose.go(c: Nat, +rows: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
143 match c:
144 case 0n:
145 Nil{}
146 case 1n+p:
147 heads(rows) <> transpose.go(p, tails(rows))
148
149def rows.zip_add(xs: List<&2, List<&2, F32>>, ys: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
150 match xs ys:
151 case Con{a, t} Con{b, u}:
152 zip_with.add(a, b) <> rows.zip_add(t, u)
153 case _ _:
154 Nil{}
155
156def rows.zip_sub(xs: List<&2, List<&2, F32>>, ys: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
157 match xs ys:
158 case Con{a, t} Con{b, u}:
159 zip_with.sub(a, b) <> rows.zip_sub(t, u)
160 case _ _:
161 Nil{}
162
163def rows.zip_mul(xs: List<&2, List<&2, F32>>, ys: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
164 match xs ys:
165 case Con{a, t} Con{b, u}:
166 zip_with.mul(a, b) <> rows.zip_mul(t, u)
167 case _ _:
168 Nil{}
169
170def rows.scale(+k: F32, xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
171 match xs:
172 case Nil{}:
173 Nil{}
174 case Con{r, t}:
175 scale.l(k, r) <> rows.scale(k, t)
176
177def rows.relu(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
178 match xs:
179 case Nil{}:
180 Nil{}
181 case Con{r, t}:
182 relu.l(r) <> rows.relu(t)
183
184# soma por colunas
185def rows.col_sums(c: Nat, +rows: List<&2, List<&2, F32>>) -> List<&2, F32>:
186 match c:
187 case 0n:
188 Nil{}
189 case 1n+p:
190 sum.go(heads(rows), 0.0) <> rows.col_sums(p, tails(rows))
191
192def rows.add_row(xs: List<&2, List<&2, F32>>, +b: List<&2, F32>) -> List<&2, List<&2, F32>>:
193 match xs:
194 case Nil{}:
195 Nil{}
196 case Con{r, t}:
197 zip_with.add(r, b) <> rows.add_row(t, b)
198
199def matmul.row(+ra: List<&2, F32>, bt: List<&2, List<&2, F32>>) -> List<&2, F32>:
200 match bt:
201 case Nil{}:
202 Nil{}
203 case Con{col, t}:
204 dot.go(ra, col, 0.0) <> matmul.row(ra, t)
205
206def matmul.rows(a: List<&2, List<&2, F32>>, +bt: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
207 match a:
208 case Nil{}:
209 Nil{}
210 case Con{r, t}:
211 matmul.row(r, bt) <> matmul.rows(t, bt)
212
213# junta todas as linhas numa lista só
214def flatten.l(rows: List<&2, List<&2, F32>>) -> List<&2, F32>:
215 match rows:
216 case Nil{}:
217 Nil{}
218 case Con{r, t}:
219 List.append(&2, F32, r, flatten.l(t))
220
221# quebra uma lista em r linhas de c números
222def chunk.l(r: Nat, +c: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
223 match r:
224 case 0n:
225 Nil{}
226 case 1n+p:
227 List.take(&2, F32, xs, c) <> chunk.l(p, c, List.drop(&2, F32, xs, c))
228
229# ---------------------------------------------------------------------
230# Mat
231# ---------------------------------------------------------------------
232
233def Mat.zeros(+r: Nat, +c: Nat) -> Mat<r, c>:
234 Mat{chunk.l(r, c, fill.l(Nat.mul(r, c), 0.0))}
235
236def Mat.fill(+r: Nat, +c: Nat, +v: F32) -> Mat<r, c>:
237 Mat{chunk.l(r, c, fill.l(Nat.mul(r, c), v))}
238
239def Mat.add(-r: Nat, -c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> Mat<r, c>:
240 match a b:
241 case Mat{x} Mat{y}:
242 Mat{rows.zip_add(x, y)}
243
244def Mat.sub(-r: Nat, -c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> Mat<r, c>:
245 match a b:
246 case Mat{x} Mat{y}:
247 Mat{rows.zip_sub(x, y)}
248
249# produto elemento a elemento (Hadamard)
250def Mat.mul(-r: Nat, -c: Nat, a: Mat<r, c>, b: Mat<r, c>) -> Mat<r, c>:
251 match a b:
252 case Mat{x} Mat{y}:
253 Mat{rows.zip_mul(x, y)}
254
255def Mat.scale(-r: Nat, -c: Nat, +k: F32, a: Mat<r, c>) -> Mat<r, c>:
256 match a:
257 case Mat{x}:
258 Mat{rows.scale(k, x)}
259
260def Mat.relu(-r: Nat, -c: Nat, a: Mat<r, c>) -> Mat<r, c>:
261 match a:
262 case Mat{x}:
263 Mat{rows.relu(x)}
264
265def Mat.transpose(-r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<c, r>:
266 match a:
267 case Mat{x}:
268 Mat{transpose.go(c, x)}
269
270# (n x k) · (k x m) = (n x m): as dimensões internas têm que ser a mesma k
271def Mat.matmul(-n: Nat, +k: Nat, +m: Nat, a: Mat<n, k>, b: Mat<k, m>) -> Mat<n, m>:
272 match a b:
273 case Mat{x} Mat{y}:
274 Mat{matmul.rows(x, transpose.go(m, y))}
275
276# (n x k) · (m x k)ᵀ = (n x m): o segundo operando já vem TRANSPOSTO (guardado como
277# m linhas de k números). Evita transpor pesos grandes a cada chamada; o tipo
278# continua exigindo a mesma dimensão interna k.
279def Mat.matmul_t(-n: Nat, -k: Nat, -m: Nat, a: Mat<n, k>, bt: Mat<m, k>) -> Mat<n, m>:
280 match a bt:
281 case Mat{x} Mat{y}:
282 Mat{matmul.rows(x, y)}
283
284# soma bias (um vetor de m números) a cada linha de uma matriz n x m
285def Mat.add_row(-n: Nat, -m: Nat, a: Mat<n, m>, b: Vec<m>) -> Mat<n, m>:
286 match a b:
287 case Mat{x} Vec{v}:
288 Mat{rows.add_row(x, v)}
289
290# soma de cada coluna: Mat<n, m> -> Vec<m>
291def Mat.col_sums(-n: Nat, +m: Nat, a: Mat<n, m>) -> Vec<m>:
292 match a:
293 case Mat{x}:
294 Vec{rows.col_sums(m, x)}
295
296# matriz · vetor: (n x k) · k = n
297def Mat.matvec(-n: Nat, +k: Nat, a: Mat<n, k>, v: Vec<k>) -> Vec<n>:
298 match a v:
299 case Mat{x} Vec{ys}:
300 Vec{heads(matmul.rows(x, ys <> Nil{}))}
301
302# Reshape: só compila com a PROVA de que o número de elementos não muda.
303def Mat.reshape(-r1: Nat, -c1: Nat, r2: Nat, +c2: Nat, -p: {Nat.mul(r1, c1) == Nat.mul(r2, c2) : Nat}, a: Mat<r1, c1>) -> Mat<r2, c2>:
304 match a:
305 case Mat{x}:
306 Mat{chunk.l(r2, c2, flatten.l(x))}
307
308# ---------------------------------------------------------------------
309# funções por linha: softmax, GELU, layernorm
310# ---------------------------------------------------------------------
311
312def max.go(xs: List<&2, F32>, cur: F32) -> F32:
313 match xs:
314 case Nil{}:
315 cur
316 case Con{+h, t}:
317 max.go(t, F32.max(cur, h))
318
319def max.l(xs: List<&2, F32>) -> F32:
320 match xs:
321 case Nil{}:
322 0.0
323 case Con{+h, t}:
324 max.go(t, h)
325
326def exp_shift(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
327 match xs:
328 case Nil{}:
329 Nil{}
330 case Con{+h, t}:
331 F32.exp((h - m : F32)) <> exp_shift(t, m)
332
333def div.l(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
334 match xs:
335 case Nil{}:
336 Nil{}
337 case Con{+h, t}:
338 (h / s : F32) <> div.l(t, s)
339
340# softmax estável: subtrai o máximo antes de exponenciar
341def softmax.fin(+es: List<&2, F32>) -> List<&2, F32>:
342 div.l(es, sum.go(es, 0.0))
343
344def softmax.l(+xs: List<&2, F32>) -> List<&2, F32>:
345 softmax.fin(exp_shift(xs, max.l(xs)))
346
347def rows.softmax(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
348 match xs:
349 case Nil{}:
350 Nil{}
351 case Con{r, t}:
352 softmax.l(r) <> rows.softmax(t)
353
354# GELU com a aproximação por tanh (a mesma do GPT-2)
355def gelu.f(+x: F32) -> F32:
356 (0.5 * (x * (1.0 + F32.tanh((0.7978846 * (x + (0.044715 * (x * (x * x : F32) : F32) : F32) : F32) : F32)) : F32) : F32) : F32)
357
358def gelu.l(xs: List<&2, F32>) -> List<&2, F32>:
359 match xs:
360 case Nil{}:
361 Nil{}
362 case Con{h, t}:
363 gelu.f(h) <> gelu.l(t)
364
365def rows.gelu(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
366 match xs:
367 case Nil{}:
368 Nil{}
369 case Con{r, t}:
370 gelu.l(r) <> rows.gelu(t)
371
372def sub_const(xs: List<&2, F32>, +k: F32) -> List<&2, F32>:
373 match xs:
374 case Nil{}:
375 Nil{}
376 case Con{+h, t}:
377 (h - k : F32) <> sub_const(t, k)
378
379def sq_sum(xs: List<&2, F32>, acc: F32) -> F32:
380 match xs:
381 case Nil{}:
382 acc
383 case Con{+h, t}:
384 sq_sum(t, (acc + (h * h : F32) : F32))
385
386# (x - média) / sqrt(variância + eps), depois * gamma + beta
387def layernorm.fin(+cs: List<&2, F32>, +d: F32, +eps: F32, +gamma: List<&2, F32>, +beta: List<&2, F32>) -> List<&2, F32>:
388 zip_with.add(zip_with.mul(scale.l((1.0 / F32.sqrt(((sq_sum(cs, 0.0) / d : F32) + eps : F32)) : F32), cs), gamma), beta)
389
390def layernorm.l(+xs: List<&2, F32>, +d: F32, +eps: F32, +gamma: List<&2, F32>, +beta: List<&2, F32>) -> List<&2, F32>:
391 layernorm.fin(sub_const(xs, (sum.go(xs, 0.0) / d : F32)), d, eps, gamma, beta)
392
393def rows.layernorm(xs: List<&2, List<&2, F32>>, +d: F32, +eps: F32, +gamma: List<&2, F32>, +beta: List<&2, F32>) -> List<&2, List<&2, F32>>:
394 match xs:
395 case Nil{}:
396 Nil{}
397 case Con{r, t}:
398 layernorm.l(r, d, eps, gamma, beta) <> rows.layernorm(t, d, eps, gamma, beta)
399
400def Mat.softmax(-r: Nat, -c: Nat, a: Mat<r, c>) -> Mat<r, c>:
401 match a:
402 case Mat{x}:
403 Mat{rows.softmax(x)}
404
405def Mat.gelu(-r: Nat, -c: Nat, a: Mat<r, c>) -> Mat<r, c>:
406 match a:
407 case Mat{x}:
408 Mat{rows.gelu(x)}
409
410# normaliza cada linha (d colunas) e aplica gamma e beta, vetores de d números
411def Mat.layernorm(-n: Nat, d: Nat, eps: F32, g: Vec<d>, bt: Vec<d>, a: Mat<n, d>) -> Mat<n, d>:
412 match g bt a:
413 case Vec{gamma} Vec{beta} Mat{x}:
414 Mat{rows.layernorm(x, F32.from_nat(d), eps, gamma, beta)}
415
416def Vec.softmax(-n: Nat, a: Vec<n>) -> Vec<n>:
417 match a:
418 case Vec{xs}:
419 Vec{softmax.l(xs)}
420
421# ---------------------------------------------------------------------
422# construção e leitura com verificação de tamanho
423# ---------------------------------------------------------------------
424
425def len.eq(xs: List<&2, F32>, n: Nat) -> Bool:
426 match xs n:
427 case Nil{} 0n:
428 True{}
429 case Con{h, t} 1n+p:
430 len.eq(t, p)
431 case _ _:
432 False{}
433
434def rows.ok(rs: List<&2, List<&2, F32>>, +c: Nat, r: Nat) -> Bool:
435 match rs r:
436 case Nil{} 0n:
437 True{}
438 case Con{row, t} 1n+p:
439 Bool.and(len.eq(row, c), rows.ok(t, c, p))
440 case _ _:
441 False{}
442
443def Vec.of2(-n: Nat, ok: Bool, xs: List<&2, F32>) -> Maybe<&2, Vec<n>>:
444 match ok:
445 case True{}:
446 Some{Vec{xs}}
447 case False{}:
448 None{}
449
450# Vec<n> a partir de uma lista, só se ela tem exatamente n números
451def Vec.of(n: Nat, +xs: List<&2, F32>) -> Maybe<&2, Vec<n>>:
452 Vec.of2(n, len.eq(xs, n), xs)
453
454def Mat.of2(-r: Nat, -c: Nat, ok: Bool, rs: List<&2, List<&2, F32>>) -> Maybe<&2, Mat<r, c>>:
455 match ok:
456 case True{}:
457 Some{Mat{rs}}
458 case False{}:
459 None{}
460
461# Mat<r, c> a partir de linhas, só se são exatamente r linhas de c números
462def Mat.of(+r: Nat, +c: Nat, +rs: List<&2, List<&2, F32>>) -> Maybe<&2, Mat<r, c>>:
463 Mat.of2(r, c, rows.ok(rs, c, r), rs)
464
465def Vec.to_list(-n: Nat, a: Vec<n>) -> List<&2, F32>:
466 match a:
467 case Vec{xs}:
468 xs
469
470def Mat.to_rows(-r: Nat, -c: Nat, a: Mat<r, c>) -> List<&2, List<&2, F32>>:
471 match a:
472 case Mat{x}:
473 x
474
475# ---------------------------------------------------------------------
476# LAWS: quando um reshape é possível, e uma prova reaproveitada do nat-lemmas
477# ---------------------------------------------------------------------
478
479# LEI (reshape_swap): r x c e c x r têm o mesmo número de elementos
480# (a multiplicação é comutativa), então esse reshape sempre existe.
481law reshape_swap:
482 for +r: Nat
483 for +c: Nat
484 {Nat.mul(r, c) == Nat.mul(c, r) : Nat}
485
486def reshape_swap(r, c):
487 NL.mul_comm(r, c)
488
489# LEI (reshape_flat): r x c tem o mesmo número de elementos que 1 x (r*c)
490# (multiplicar por 1 não muda nada), então achatar uma matriz sempre existe.
491law reshape_flat:
492 for +r: Nat
493 for +c: Nat
494 {Nat.mul(r, c) == Nat.mul(1n, Nat.mul(r, c)) : Nat}
495
496def reshape_flat(r, c):
497 Equal.sym(Nat, Nat.mul(1n, Nat.mul(r, c)), Nat.mul(r, c), NL.mul_one_l(Nat.mul(r, c)))
498
499# Mat<r, c> -> Mat<1, r*c>: a prova vem da lei acima
500def Mat.flatten(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<1n, Nat.mul(r, c)>:
501 Mat.reshape(r, c, 1n, Nat.mul(r, c), reshape_flat(r, c), a)
502
503# Mat<r, c> -> Mat<c, r> pelos mesmos dados (não é a transposta!): só reinterpreta
504def Mat.reshape_swap(+r: Nat, +c: Nat, a: Mat<r, c>) -> Mat<c, r>:
505 Mat.reshape(r, c, c, r, reshape_swap(r, c), a)