spec.bend
bendtorch@1.0.0.0 · 128 lines · 0 definitions
1# The mathematical model the laws speak about: vectors and matrices as lists of Nat, a matrix as its list of rows.
2# Nothing here runs in a model; it is the specification that LAWS.bend states facts about. Nat makes the algebra
3# exact (no rounding), so a law proved here says the formulas are right; whether the F32 code computes them to F32's
4# precision is tested against PyTorch, never proved (docs/design.md, "Proved, tested, trusted").
5
6import Base
7import bendtorch@1.0.0.0/main.bend as T
8
9# the dot product of the common prefix
10def dotn(xs: List<&2, Nat>, ys: List<&2, Nat>) -> Nat:
11 match xs ys:
12 case Con{x, xt} Con{y, yt}:
13 Nat.add(Nat.mul(x, y), dotn(xt, yt))
14 case _ _:
15 0n
16
17# vector sum; the shorter vector counts as padded with zeros
18def vadd(xs: List<&2, Nat>, ys: List<&2, Nat>) -> List<&2, Nat>:
19 match xs ys:
20 case Nil{} _:
21 ys
22 case Con{x, xt} Nil{}:
23 Con{x, xt}
24 case Con{x, xt} Con{y, yt}:
25 Nat.add(x, y) <> vadd(xt, yt)
26
27def vscale(+k: Nat, xs: List<&2, Nat>) -> List<&2, Nat>:
28 match xs:
29 case Nil{}:
30 Nil{}
31 case Con{x, t}:
32 Nat.mul(k, x) <> vscale(k, t)
33
34# M·v: one dot product per row (a linear layer's forward pass, y = W·x)
35def mv(rows: List<&2, List<&2, Nat>>, +v: List<&2, Nat>) -> List<&2, Nat>:
36 match rows:
37 case Nil{}:
38 Nil{}
39 case Con{r, t}:
40 dotn(r, v) <> mv(t, v)
41
42# Mᵀ·g: the rows weighted by g and summed (a linear layer's backward pass to its input, dx = Wᵀ·dy)
43def mtv(rows: List<&2, List<&2, Nat>>, g: List<&2, Nat>) -> List<&2, Nat>:
44 match rows g:
45 case Con{r, rt} Con{gi, gt}:
46 vadd(vscale(gi, r), mtv(rt, gt))
47 case _ _:
48 Nil{}
49
50# g ⊗ v: row i is g[i]·v (a linear layer's weight gradient, dW = dy ⊗ x)
51def outer(g: List<&2, Nat>, +v: List<&2, Nat>) -> List<&2, List<&2, Nat>>:
52 match g:
53 case Nil{}:
54 Nil{}
55 case Con{gi, gt}:
56 vscale(gi, v) <> outer(gt, v)
57
58# the sum of the products of matching entries of two matrices (Frobenius inner product)
59def frob(a: List<&2, List<&2, Nat>>, b: List<&2, List<&2, Nat>>) -> Nat:
60 match a b:
61 case Con{ar, at} Con{br, bt}:
62 Nat.add(dotn(ar, br), frob(at, bt))
63 case _ _:
64 0n
65
66# a <= b, by recursion (simpler to reason about than Nat.is_le, which goes through Nat.cmp)
67def le(a: Nat, b: Nat) -> Bool:
68 match a b:
69 case 0n _:
70 True{}
71 case 1n+p 0n:
72 False{}
73 case 1n+p 1n+q:
74 le(p, q)
75
76# ---- linear layers of the model as BendTorch layers, so the laws can run T.Seq and T.Fork's own code on them ----
77
78# row-wise sum of two matrices (the shorter counts as padded)
79def madd(a: List<&2, List<&2, Nat>>, b: List<&2, List<&2, Nat>>) -> List<&2, List<&2, Nat>>:
80 match a b:
81 case Nil{} _:
82 b
83 case Con{r, rt} Nil{}:
84 Con{r, rt}
85 case Con{r, rt} Con{s, st}:
86 vadd(r, s) <> madd(rt, st)
87
88# y = W·x, saving x
89def lin_fwd(+w: List<&2, List<&2, Nat>>, +x: List<&2, Nat>) -> T.Fwd<List<&2, List<&2, Nat>>, List<&2, Nat>, List<&2, Nat>>:
90 T.Fwd{w, mv(w, x), x}
91
92# dx = Wᵀ·dy, and dW += dy ⊗ x
93def lin_bwd(+w: List<&2, List<&2, Nat>>, +x: List<&2, Nat>, +dy: List<&2, Nat>, g: List<&2, List<&2, Nat>>) -> T.Bwd<List<&2, List<&2, Nat>>, List<&2, Nat>>:
94 T.Bwd{w, mtv(w, dy), madd(g, outer(dy, x))}
95
96def LinL() -> T.Layer<List<&2, List<&2, Nat>>, List<&2, Nat>, List<&2, Nat>, List<&2, Nat>>:
97 T.Layer{p => x => lin_fwd(p, x), p => s => dy => g => lin_bwd(p, s, dy, g), op => p => g => st => T.no_step(List<&2, List<&2, Nat>>, op, p, g, st), u => Nil{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
98
99# two linear layers in sequence, and side by side on one input, through BendTorch's combinators
100def SeqL() -> T.Layer<T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>>:
101 T.Seq(~List<&2, List<&2, Nat>>, ~List<&2, List<&2, Nat>>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~LinL(), ~LinL())
102
103def dupv(+x: List<&2, Nat>) -> T.Two<List<&2, Nat>, List<&2, Nat>>:
104 T.Two{x, x}
105
106def ForkL() -> T.Layer<T.Two<List<&2, List<&2, Nat>>, List<&2, List<&2, Nat>>>, List<&2, Nat>, T.Two<List<&2, Nat>, List<&2, Nat>>, T.Two<List<&2, Nat>, List<&2, Nat>>>:
107 T.Fork(~List<&2, List<&2, Nat>>, ~List<&2, List<&2, Nat>>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~List<&2, Nat>, ~(x => dupv(x)), ~(a => b => vadd(a, b)), ~LinL(), ~LinL())
108
109def fy(-P: Type, -Y: Type, -S: Type, r: T.Fwd<P, Y, S>) -> Y:
110 match r:
111 case T.Fwd{p, y, s}:
112 y
113
114def fsv(-P: Type, -Y: Type, -S: Type, r: T.Fwd<P, Y, S>) -> S:
115 match r:
116 case T.Fwd{p, y, s}:
117 s
118
119def bdx(-P: Type, -X: Type, r: T.Bwd<P, X>) -> X:
120 match r:
121 case T.Bwd{p, dx, g}:
122 dx
123
124# the dot product of two pairs of vectors
125def pdot(a: T.Two<List<&2, Nat>, List<&2, Nat>>, b: T.Two<List<&2, Nat>, List<&2, Nat>>) -> Nat:
126 match a b:
127 case T.Two{a1, a2} T.Two{b1, b2}:
128 Nat.add(dotn(a1, b1), dotn(a2, b2))