main.bend

bend-ml-tensor@0.1.2.0 · 505 lines · 78 definitions

1# bend-ml-tensor: tensors with the shape in the TYPE. A shape error is a type error.
2#
3#   import bend-ml-tensor@0.1.2.0/main.bend as T
4#
5# Vec<n> is a vector of n F32 numbers; Mat<r, c> is an r x c matrix (a list of r
6# rows with c numbers each). The dimensions are erased type parameters: they cost
7# nothing at run time, but the checker verifies them in every operation.
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# lists of F32 (internal; the public operations fix the dimensions)
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# lists of rows (internal)
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# transpose of a matrix with `c` columns (c is the one that shrinks, so it comes first)
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# column sums
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# joins all the rows into a single list
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# splits a list into r rows of c numbers
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# element-wise product (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): the inner dimensions must be the same 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): the second operand arrives already TRANSPOSED (stored as
277# m rows of k numbers). It avoids transposing large weights on every call; the type
278# still requires the same inner dimension 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# adds a bias (a vector of m numbers) to each row of an n x m matrix
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# sum of each column: 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# matrix · vector: (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: compiles only with the PROOF that the number of elements does not change.
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# per-row functions: 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# stable softmax: subtracts the maximum before exponentiating
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 with the tanh approximation (the same as 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 - mean) / sqrt(variance + eps), then * 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# normalizes each row (d columns) and applies gamma and beta, vectors of d numbers
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# construction and reading with size checking
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> from a list, only if it has exactly n numbers
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> from rows, only if they are exactly r rows of c numbers
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: when a reshape is possible, and a proof reused from nat-lemmas
477# ---------------------------------------------------------------------
478
479# LAW (reshape_swap): r x c and c x r have the same number of elements
480# (multiplication is commutative), so this reshape always exists.
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# LAW (reshape_flat): r x c has the same number of elements as 1 x (r*c)
490# (multiplying by 1 changes nothing), so flattening a matrix always exists.
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>: the proof comes from the law above
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> over the same data (it is NOT the transpose!): it only reinterprets
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)