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	
6	import Base
7	import ./main.bend as T
8	
9	# the dot product of the common prefix
10	def 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
18	def 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	
27	def 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)
35	def 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)
43	def 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)
51	def 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)
59	def 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)
67	def 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)
79	def 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
89	def 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
93	def 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	
96	def 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
100	def 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	
103	def dupv(+x: List<&2, Nat>) -> T.Two<List<&2, Nat>, List<&2, Nat>>:
104	  T.Two{x, x}
105	
106	def 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	
109	def 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	
114	def 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	
119	def 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
125	def 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))
129	
130	# ---- The order of a row's additions (the law block_sum). A sum here is the tree of additions it makes over the cells'
131	# positions, not a number: two ways of summing that build the same tree add the same pairs in the same order, so they
132	# give the same bits in F32, or in any arithmetic, with nothing assumed about rounding. ----
133	
134	type Sum is Data:
135	  SCell{i: Nat}
136	  SAdd{a: Sum, b: Sum}
137	
138	# 2^p
139	def pw(p: Nat) -> Nat:
140	  match p:
141	    case 0n:
142	      1n
143	    case 1n+q:
144	      +h = pw(q)
145	      Nat.add(h, h)
146	
147	# the 2^p cells from cell i summed as a balanced tree of pairs: how a GPU leaf sums its block (the device's MSE and
148	# cross-entropy block sums)
149	def bal(+p: Nat, +i: Nat) -> Sum:
150	  match p:
151	    case 0n:
152	      SCell{i}
153	    case 1n+q:
154	      SAdd{bal(q, i), bal(q, Nat.add(i, pw(q)))}
155	
156	# k blocks of 2^p cells, the first from cell lo, each summed by bal; with p = 0, a row of k cells
157	def blocks(+p: Nat, k: Nat, +lo: Nat) -> List<&2, Sum>:
158	  match k:
159	    case 0n:
160	      Nil{}
161	    case 1n+j:
162	      bal(p, lo) <> blocks(p, j, Nat.add(lo, pw(p)))
163	
164	# one pass of main.bend's tree_sum (ts_pass): neighbours added in pairs, an odd last cell carried up as it is
165	def spass(xs: List<&2, Sum>) -> List<&2, Sum>:
166	  match xs:
167	    case Nil{}:
168	      Nil{}
169	    case Con{a, Nil{}}:
170	      a <> Nil{}
171	    case Con{a, Con{b, t}}:
172	      SAdd{a, b} <> spass(t)
173	
174	# r passes of tree_sum; tree_sum (ts_go) stops once one cell is left, and a pass leaves one cell as it is, so the sum of
175	# m cells is passes(r, row) for any r at least the number of passes it makes
176	def passes(r: Nat, xs: List<&2, Sum>) -> List<&2, Sum>:
177	  match r:
178	    case 0n:
179	      xs
180	    case 1n+q:
181	      spass(passes(q, xs))