1	# BendTorch's layers over bend-ml-tensor-fork: the same layers, types and bits as torch/main.bend's, faster where the
2	# work splits into parallel blocks.
3	#
4	#   import ../torch/main.bend as T
5	#   import ../torch/fast.bend as F
6	#   F.LinearFirstF(100n, 784n, 128n, 3n)     # instead of T.LinearFirst(100n, 784n, 128n, 3n)
7	#
8	# LinearF, LinearFirstF and LinearNBF are main.bend's Linear, LinearFirst and LinearNB with their three products
9	# (x·Wᵀ, dY·W, dYᵀ·X added into dW) taken from bend-ml-tensor-fork instead of bend-ml-tensor-array, and their optimizer
10	# step in blocks. SiLU, ActIn, Mul2, Add2, Residual, Mat.dup and Mat.sum2 are main.bend's with their element-wise loops
11	# in blocks (tensor-fork's par1..par4), the optimizer step (mat_step, lin_step) too. Every number is computed by
12	# main.bend's own loop or tensor-array's own cells, so the bits are main.bend's at every depth and thread count (tests:
13	# demo/fast_test.bend). Small matrices run in one block (ew_depth), as in main.bend.
14	#
15	# Outside the proofs, on purpose: tensor-fork shares arrays through Base's @unsafe Array.fork, so `bend` names the defs
16	# below that reach it, and no file with laws imports this one (`make laws` checks both). All of it is Bend.
17	
18	import Base
19	import bend-ml-tensor-array@0.1.14.0/main.bend as TA
20	import bend-ml-tensor-fork@0.1.2.0/main.bend as TF
21	import bendtorch@1.2.0.0/main.bend as T
22	
23	# =====================================================================================================
24	# Element-wise work and the optimizer step in parallel blocks (bend-ml-tensor-fork's par1..par4): every number is
25	# computed by main.bend's own loop, so the bits are main.bend's at every depth and thread count. Small matrices run in
26	# one block (ew_depth).
27	# =====================================================================================================
28	
29	# log2 of the blocks for cnt numbers: one block below 131072 (MNIST's and the tiny Llama's matrices; at 100352 the
30	# blocks cost MNIST more than they gave on 16 threads), else one per 16384, at most 64
31	def ew_d2(small: Bool, big: Bool, +cnt: Nat) -> Nat:
32	  match small big:
33	    case True{} _:
34	      0n
35	    case False{} True{}:
36	      6n
37	    case False{} False{}:
38	      TF.depth(Nat.div(cnt, 16384n))
39	
40	def ew_depth(+cnt: Nat) -> Nat:
41	  ew_d2(Nat.is_lt(cnt, 131072n), Nat.is_le(1048576n, cnt), cnt)
42	
43	# ---- the optimizer step on a matrix ----
44	
45	def p2_b2(q: TA.P2) -> TF.B2<F32, F32>:
46	  match q:
47	    case TA.P2{a, b}:
48	      TF.B2{a, b}
49	
50	def sgd_leaf(+lr: F32, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
51	  match r:
52	    case TF.B2{w, g}:
53	      p2_b2(T.sgd_l(cnt, i0, lr, Array.get(F32, w, i0), Array.get(F32, g, i0)))
54	
55	def mom_b(q: T.Q3) -> TF.B3<F32, F32, F32>:
56	  match q:
57	    case T.Q3{w, g, m}: TF.B3{w, g, m}
58	
59	def mom_leaf(+op: T.Opt, +i0: U32, +cnt: Nat, r: TF.B3<F32, F32, F32>) -> TF.B3<F32, F32, F32>:
60	  match op r:
61	    case T.Momentum{+lr, +mu} TF.B3{w, g, m}:
62	      mom_b(T.mom_l(cnt, i0, lr, mu, Array.get(F32, w, i0), Array.get(F32, g, i0), Array.get(F32, m, i0)))
63	    case _ r: r
64	
65	def q4_b(q: T.Q4) -> TF.B4<F32, F32, F32, F32>:
66	  match q:
67	    case T.Q4{w, g, m, v}: TF.B4{w, g, m, v}
68	
69	def opt_leaf(+o: T.Opt, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, F32, F32>) -> TF.B4<F32, F32, F32, F32>:
70	  match o r:
71	    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TF.B4{w, g, m, v}:
72	      q4_b(T.adam_l(cnt, i0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, i0), Array.get(F32, g, i0), Array.get(F32, m, i0), Array.get(F32, v, i0)))
73	    case _ r: r
74	
75	def sgd_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> T.MatStep<r, c>:
76	  match q:
77	    case TF.B2{w, g}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.NoMom{}}
78	
79	def mom_fin(-r: Nat, -c: Nat, q: TF.B3<F32, F32, F32>) -> T.MatStep<r, c>:
80	  match q:
81	    case TF.B3{w, g, m}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.OneMom{TA.Mat{m}}}
82	
83	def opt_fin(-r: Nat, -c: Nat, q: TF.B4<F32, F32, F32, F32>) -> T.MatStep<r, c>:
84	  match q:
85	    case TF.B4{w, g, m, v}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.TwoMom{TA.Mat{m}, TA.Mat{v}}}
86	
87	# ---- AdamW over packed BF16 moments (CAdamWBF16), in blocks of whole words ----
88	# Two moments share a word, so the blocks split the words, not the elements: a block of words i0 .. i0 + cnt - 1 runs
89	# main.bend's own loop (adamh_l, or adamhv_l when v stays F32) over elements 2·i0 .. 2·(i0 + cnt) - 1, the last block
90	# stopping at the matrix's n. No word is written by two blocks, and each element's arithmetic is main.bend's.
91	
92	type PE is Data:
93	  PE{o: T.Opt, n: U32}
94	
95	# the elements of a block of cnt words from element e0, the last block stopping at n
96	def ph_cnt(+n: U32, +e0: U32, +cnt: Nat) -> Nat:
97	  +k = U32.from_nat(Nat.add(cnt, cnt))
98	  U32.to_nat(T.ph_pick(U32.is_lt(k, (n - e0 : U32)), k, (n - e0 : U32)))
99	
100	def pq_b(q: T.PQ) -> TF.B4<F32, F32, U32, U32>:
101	  match q:
102	    case T.PQ{w, g, m, v}: TF.B4{w, g, m, v}
103	
104	def pqv_b(q: T.PQV) -> TF.B4<F32, F32, U32, F32>:
105	  match q:
106	    case T.PQV{w, g, m, v}: TF.B4{w, g, m, v}
107	
108	def ph_go(+o: T.Opt, +n: U32, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, U32>) -> TF.B4<F32, F32, U32, U32>:
109	  match o r:
110	    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TF.B4{w, g, m, v}:
111	      +e0 = (i0 + i0 : U32)
112	      pq_b(T.adamh_l(ph_cnt(n, e0, cnt), e0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, e0), Array.get(F32, g, e0), Array.get(U32, m, i0), Array.get(U32, v, i0)))
113	    case _ r: r
114	
115	def ph_leaf(+e: PE, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, U32>) -> TF.B4<F32, F32, U32, U32>:
116	  match e:
117	    case PE{+o, +n}: ph_go(o, n, i0, cnt, r)
118	
119	def phv_go(+o: T.Opt, +n: U32, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, F32>) -> TF.B4<F32, F32, U32, F32>:
120	  match o r:
121	    case T.Adam{+lr, +b1, +b2, +eps, +wd, +dec, +step, +bc2, +c1, +c2} TF.B4{w, g, m, v}:
122	      +e0 = (i0 + i0 : U32)
123	      pqv_b(T.adamhv_l(ph_cnt(n, e0, cnt), e0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, w, e0), Array.get(F32, g, e0), Array.get(U32, m, i0), Array.get(F32, v, e0)))
124	    case _ r: r
125	
126	def phv_leaf(+e: PE, +i0: U32, +cnt: Nat, r: TF.B4<F32, F32, U32, F32>) -> TF.B4<F32, F32, U32, F32>:
127	  match e:
128	    case PE{+o, +n}: phv_go(o, n, i0, cnt, r)
129	
130	def ph_fin(-r: Nat, -c: Nat, +n: U32, q: TF.B4<F32, F32, U32, U32>) -> T.MatStep<r, c>:
131	  match q:
132	    case TF.B4{w, g, m, v}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.Packed{T.PLeaf{n, m, v}}}
133	
134	def phv_fin(-r: Nat, -c: Nat, +n: U32, q: TF.B4<F32, F32, U32, F32>) -> T.MatStep<r, c>:
135	  match q:
136	    case TF.B4{w, g, m, v}: T.MatStep{TA.Mat{w}, TA.Mat{g}, T.Packed{T.PLeafV{n, m, v}}}
137	
138	# the words of r·c elements, ceil(r·c / 2)
139	def ph_nw(+r: Nat, +c: Nat) -> Nat:
140	  Nat.div(Nat.add(Nat.mul(r, c), 1n), 2n)
141	
142	def ph_leafs(+r: Nat, +c: Nat, +o: T.Opt, wa: Array<F32>, ga: Array<F32>, t: T.PTree) -> T.MatStep<r, c>:
143	  match o t:
144	    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} T.PLeaf{+n, ma, va}:
145	      ph_fin(r, c, n, TF.par4(~PE, ~F32, ~F32, ~U32, ~U32, ~ph_leaf, PE{o, n}, ew_depth(Nat.mul(r, c)), ph_nw(r, c), wa, ga, ma, va))
146	    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} T.PLeafV{+n, ma, va}:
147	      phv_fin(r, c, n, TF.par4(~PE, ~F32, ~F32, ~U32, ~F32, ~phv_leaf, PE{o, n}, ew_depth(Nat.mul(r, c)), ph_nw(r, c), wa, ga, ma, va))
148	    case _ t:
149	      T.MatStep{TA.Mat{wa}, TA.Mat{ga}, T.Packed{t}}
150	
151	def ph_mat(+r: Nat, +c: Nat, +o: T.Opt, wa: Array<F32>, ga: Array<F32>, t: T.PTree) -> T.MatStep<r, c>:
152	  match t:
153	    case T.PNew{v16}:
154	      ph_leafs(r, c, o, wa, ga, T.ph_new(Nat.mul(r, c), v16))
155	    case t:
156	      ph_leafs(r, c, o, wa, ga, t)
157	
158	# SGD splits two arrays, momentum three, Adam four. No unused parameter-sized buffers.
159	def opt_mat(+r: Nat, +c: Nat, +o: T.Opt, w: TA.Mat<r, c>, g: TA.Mat<r, c>, st: T.OptState<TA.Mat<r, c>>) -> T.MatStep<r, c>:
160	  match o w g st:
161	    case T.SGD{+lr} TA.Mat{wd} TA.Mat{gd} _:
162	      sgd_fin(r, c, TF.par2(~F32, ~F32, ~F32, ~sgd_leaf, lr, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd))
163	    case T.Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} T.OneMom{TA.Mat{md}}:
164	      mom_fin(r, c, TF.par3(~T.Opt, ~F32, ~F32, ~F32, ~mom_leaf, o, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd, md))
165	    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} TA.Mat{wa} TA.Mat{ga} T.TwoMom{TA.Mat{ma}, TA.Mat{va}}:
166	      opt_fin(r, c, TF.par4(~T.Opt, ~F32, ~F32, ~F32, ~F32, ~opt_leaf, o, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wa, ga, ma, va))
167	    # packed BF16 moments (CAdamWBF16): blocks of whole words
168	    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} TA.Mat{wa} TA.Mat{ga} T.Packed{t}:
169	      ph_mat(r, c, o, wa, ga, t)
170	    case _ w g st: T.MatStep{w, g, st}
171	
172	def mat_step(+r: Nat, +c: Nat, +op: T.Opt, p: TA.Mat<r, c>, g: TA.Mat<r, c>, st: T.OptState<TA.Mat<r, c>>) -> T.Step<TA.Mat<r, c>>:
173	  T.mat_step_fin(r, c, opt_mat(r, c, op, p, g, st))
174	
175	def lin_step2(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin<i, o>, g: T.Lin<i, o>, st: T.Two<T.OptState<TA.Mat<o, i>>, T.OptState<TA.Mat<1n, o>>>) -> T.Step<T.Lin<i, o>>:
176	  match p g st:
177	    case T.Lin{w, b} T.Lin{gw, gb} T.Two{sw, sb}:
178	      T.lin_s2(i, o, opt_mat(o, i, op, w, gw, sw), opt_mat(1n, o, op, b, gb, sb))
179	
180	def lin_step(+i: Nat, +o: Nat, +op: T.Opt, p: T.Lin<i, o>, g: T.Lin<i, o>, st: T.OptState<T.Lin<i, o>>) -> T.Step<T.Lin<i, o>>:
181	  lin_step2(i, o, op, p, g, T.OptState.split(TA.Mat<o, i>, TA.Mat<1n, o>, T.OptState.map(T.Lin<i, o>, T.Two<TA.Mat<o, i>, TA.Mat<1n, o>>, st, x => T.lin_pair(i, o, x), x => T.lin_pair(i, o, x))))
182	
183	# ---- element-wise maps, zips and copies ----
184	
185	def ew1_leaf(~f: @+a: F32 -> F32, +e: Unit, +i0: U32, +cnt: Nat, a: Array<F32>) -> Array<F32>:
186	  T.ew1_l(~f, cnt, i0, Array.get(F32, a, i0))
187	
188	def map(~f: @+a: F32 -> F32, +r: Nat, +c: Nat, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
189	  match m:
190	    case TA.Mat{d}:
191	      TA.Mat{TF.par1(~Unit, ~F32, ~(e => i0 => k => a => ew1_leaf(~f, e, i0, k, a)), Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), d)}
192	
193	
194	def ew2_leaf(~f: @+a: F32 -> @+b: F32 -> F32, +e: Unit, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
195	  match r:
196	    case TF.B2{a, b}:
197	      p2_b2(T.ew2_l(~f, cnt, i0, Array.get(F32, a, i0), Array.get(F32, b, i0)))
198	
199	def zip_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> TA.MPair<r, c>:
200	  match q:
201	    case TF.B2{a, b}:
202	      TA.MPair{TA.Mat{a}, TA.Mat{b}}
203	
204	# a[i] = f(a[i], b[i]); b comes back unchanged
205	def zip(~f: @+a: F32 -> @+b: F32 -> F32, +r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.MPair<r, c>:
206	  match a b:
207	    case TA.Mat{x} TA.Mat{y}:
208	      zip_fin(r, c, TF.par2(~Unit, ~F32, ~F32, ~(e => i0 => k => q => ew2_leaf(~f, e, i0, k, q)), Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), x, y))
209	
210	def a2_b2(q: T.A2) -> TF.B2<F32, F32>:
211	  match q:
212	    case T.A2{a, b}:
213	      TF.B2{a, b}
214	
215	# numbers i0 .. i0+cnt-1 of a into b
216	def copy_leaf(+e: Unit, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
217	  match cnt r:
218	    case 0n r:
219	      r
220	    case 1n+k TF.B2{a, b}:
221	      a2_b2(T.copy_n(k, i0, i0, b, Array.get(F32, a, i0)))
222	
223	def dup_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
224	  match q:
225	    case TF.B2{a, b}:
226	      T.Two{TA.Mat{a}, TA.Mat{b}}
227	
228	# two copies of a matrix (T.Mat.dup)
229	def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
230	  match m:
231	    case TA.Mat{a}:
232	      dup_fin(r, c, TF.par2(~Unit, ~F32, ~F32, ~copy_leaf, Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), a, T.zeros_a(Nat.mul(r, c))))
233	
234	def mpair_a(-r: Nat, -c: Nat, q: TA.MPair<r, c>) -> TA.Mat<r, c>:
235	  match q:
236	    case TA.MPair{a, b}:
237	      a
238	
239	# a + b, dropping b (T.Mat.sum2)
240	def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
241	  mpair_a(r, c, zip(~T.f32_add, r, c, a, b))
242	
243	# ---- Linear: y = x·Wᵀ + b ----
244	
245	def lin_fwd(+n: Nat, +i: Nat, +o: Nat, par: Nat, p: T.Lin<i, o>, x: TA.Mat<n, i>) -> T.Fwd<T.Lin<i, o>, TA.Mat<n, o>, TA.Mat<n, i>>:
246	  match p:
247	    case T.Lin{w, b}:
248	      T.lin_f2(n, i, o, b, TF.matmul_nt(n, i, o, par, x, w))
249	
250	def lin_b2(+n: Nat, +i: Nat, +o: Nat, par: Nat, b: TA.Mat<1n, o>, x: TA.Mat<n, i>, g: T.Lin<i, o>, r: TA.MMul<n, o, i>) -> T.Bwd<T.Lin<i, o>, TA.Mat<n, i>>:
251	  match g r:
252	    case T.Lin{gw, gb} TA.MMul{dy, w, dx}:
253	      T.lin_b3(n, i, o, w, b, dx, gb, TF.matmul_tn_acc(o, n, i, par, dy, x, gw))
254	
255	def lin_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: T.Lin<i, o>, x: TA.Mat<n, i>, dy: TA.Mat<n, o>, g: T.Lin<i, o>) -> T.Bwd<T.Lin<i, o>, TA.Mat<n, i>>:
256	  match p:
257	    case T.Lin{w, b}:
258	      lin_b2(n, i, o, par, b, x, g, TF.matmul_precise(n, o, i, par, dy, w))
259	
260	def LinearF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
261	  T.Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => lin_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => T.lin_zeros(i, o), p => T.lin_dump(i, o, p), p => xs => off => T.lin_load(i, o, p, xs, off), p => xs => off => T.lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))}
262	
263	# ---- LinearFirst: no dX (the first layer's input needs no gradient) ----
264	
265	def linf_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: T.Lin<i, o>, x: TA.Mat<n, i>, dy: TA.Mat<n, o>, g: T.Lin<i, o>) -> T.Bwd<T.Lin<i, o>, TA.Mat<n, i>>:
266	  match p g:
267	    case T.Lin{w, b} T.Lin{gw, gb}:
268	      T.linf_b3(n, i, o, w, b, gb, TF.matmul_tn_acc(o, n, i, par, dy, x, gw))
269	
270	def LinearFirstF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
271	  T.Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => linf_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => T.lin_zeros(i, o), p => T.lin_dump(i, o, p), p => xs => off => T.lin_load(i, o, p, xs, off), p => xs => off => T.lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))}
272	
273	# ---- LinearNB: y = x·Wᵀ, no bias (Llama's projections) ----
274	
275	def lnb_b2(+n: Nat, +i: Nat, +o: Nat, +par: Nat, x: TA.Mat<n, i>, g: TA.Mat<o, i>, r: TA.MMul<n, o, i>) -> T.Bwd<TA.Mat<o, i>, TA.Mat<n, i>>:
276	  match r:
277	    case TA.MMul{dy, w, dx}:
278	      T.lnb_b3(n, i, o, w, dx, TF.matmul_tn_acc(o, n, i, par, dy, x, g))
279	
280	def LinearNBF(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<TA.Mat<o, i>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
281	  T.Layer{w => x => T.lnb_f2(n, i, o, TF.matmul_nt(n, i, o, par, x, w)), w => s => dy => g => lnb_b2(n, i, o, par, s, g, TF.matmul_precise(n, o, i, par, dy, w)), op => p => g => st => mat_step(o, i, op, p, g, st), u => TA.Mat.zeros(o, i), p => T.Mat.dump(o, i, p), p => xs => off => T.Mat.restore(o, i, p, xs, off), p => xs => off => T.Mat.store(o, i, p, xs, off), u => U32.from_nat(Nat.mul(o, i))}
282	
283	# ---- element-wise layers ----
284	
285	def acti_f3(~f: @+a: F32 -> F32, +n: Nat, +c: Nat, m: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, TA.Mat<n, c>>:
286	  match m:
287	    case T.Two{x, s}:
288	      T.Fwd{Unit{}, map(~f, n, c, x), s}
289	
290	# T.ActIn: f on a copy, the input saved
291	def ActIn(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> T.Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
292	  T.Layer{p => x => acti_f3(~f, n, c, Mat.dup(n, c, x)), p => s => dy => g => T.act_b2(n, c, zip(~d, n, c, dy, s)), op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
293	
294	def SiLU(+n: Nat, +c: Nat) -> T.Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
295	  ActIn(~T.silu_v, ~T.silu_d, n, c)
296	
297	def add2_f(+n: Nat, +c: Nat, x: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, Unit>:
298	  match x:
299	    case T.Two{a, b}:
300	      T.Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}}
301	
302	def Add2(+n: Nat, +c: Nat) -> T.Layer<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Unit>:
303	  T.Layer{p => x => add2_f(n, c, x), p => s => dy => g => T.Bwd{p, Mat.dup(n, c, dy), g}, op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
304	
305	def mul2_f2(+n: Nat, +c: Nat, da: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, db: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
306	  match da db:
307	    case T.Two{a1, a2} T.Two{b1, b2}:
308	      T.mul2_f3(n, c, a2, b2, zip(~T.f32_mul, n, c, a1, b1))
309	
310	def mul2_f(+n: Nat, +c: Nat, x: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Fwd<Unit, TA.Mat<n, c>, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
311	  match x:
312	    case T.Two{a, b}:
313	      mul2_f2(n, c, Mat.dup(n, c, a), Mat.dup(n, c, b))
314	
315	def mul2_b2(+n: Nat, +c: Nat, s: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, d: T.Two<TA.Mat<n, c>, TA.Mat<n, c>>) -> T.Bwd<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
316	  match s d:
317	    case T.Two{a, b} T.Two{d1, d2}:
318	      T.Bwd{Unit{}, T.Two{mpair_a(n, c, zip(~T.f32_mul, n, c, d1, b)), mpair_a(n, c, zip(~T.f32_mul, n, c, d2, a))}, Unit{}}
319	
320	def Mul2(+n: Nat, +c: Nat) -> T.Layer<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>>:
321	  T.Layer{p => x => mul2_f(n, c, x), p => s => dy => g => mul2_b2(n, c, s, Mat.dup(n, c, dy)), op => p => g => st => T.no_step(Unit, op, p, g, st), u => Unit{}, p => T.PL{p, Nil{}}, p => xs => off => T.PA{p, xs, off}, p => xs => off => T.PA{p, xs, off}, u => 0}
322	
323	# T.Residual with these copies and sums (T.ResidualOf, the composition the law residual is about)
324	def Residual(~n: Nat, ~c: Nat, ~P: Type, ~S: Type, ~f: T.Layer<P, TA.Mat<n, c>, TA.Mat<n, c>, S>) -> T.Layer<T.Two<T.Two<Unit, P>, Unit>, TA.Mat<n, c>, TA.Mat<n, c>, T.Two<T.Two<Unit, S>, Unit>>:
325	  T.ResidualOf(~TA.Mat<n, c>, ~P, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~Add2(n, c), ~f)
326	
327	# main.bend's RMSNorm, RoPE, Attention, MSE and CrossEntropy under their own names (Embedding below, with its step in
328	# blocks), so that a program that takes its layers from here can take device versions of the same names from another
329	# file by its import line alone (they run as in main.bend)
330	def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> T.Layer<TA.Mat<1n, d>, TA.Mat<n, d>, TA.Mat<n, d>, TA.Mat<n, d>>:
331	  T.RMSNorm(n, d, eps)
332	
333	def RoPE(+n: Nat, +w: Nat, +t: Nat, +h: Nat, +half: Nat, +base: F32, -e: {w == Nat.mul(h, Nat.add(half, half)) : Nat}) -> T.Layer<Unit, TA.Mat<n, w>, TA.Mat<n, w>, Unit>:
334	  T.RoPE(n, w, t, h, half, base, e)
335	
336	def Attention(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +t: Nat, +kv: Nat, +g: Nat, +hd: Nat, +sc: F32, -e1: {n == Nat.mul(b, t) : Nat}, -e2: {w == Nat.mul(Nat.mul(kv, g), hd) : Nat}, -e3: {kvw == Nat.mul(kv, hd) : Nat}) -> T.Layer<Unit, T.Two<TA.Mat<n, w>, T.Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>, TA.Mat<n, w>, T.AttnS>:
337	  T.Attention(n, w, kvw, b, t, kv, g, hd, sc, e1, e2, e3)
338	
339	def MSE(+n: Nat, +c: Nat, y: TA.Mat<n, c>, t: TA.Mat<n, c>) -> T.LossOut<TA.Mat<n, c>>:
340	  T.MSE(n, c, y, t)
341	
342	# main.bend's Embedding with its optimizer step in blocks (mat_step): the table is SmolLM2's largest matrix (49152 x 576,
343	# tied to the output head), and main.bend's step updates it on one thread
344	def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
345	  T.Layer{p => x => T.emb_f(n, v, d, x, p), p => s => dy => g => T.emb_bwd(n, v, d, p, s, dy, g), op => p => g => st => mat_step(v, d, op, p, g, st), u => TA.Mat.zeros(v, d), p => T.Mat.dump(v, d, p), p => xs => off => T.Mat.restore(v, d, p, xs, off), p => xs => off => T.Mat.store(v, d, p, xs, off), u => U32.from_nat(Nat.mul(v, d))}
346	
347	def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat<n, c>, +labels: List<&2, Nat>) -> T.LossOut<TA.Mat<n, c>>:
348	  T.CrossEntropy(n, c, z, labels)
349	
350	# ---- main.bend's names for the same layers (#13) ----
351	#
352	# A program written against torch/main.bend switches to this lane by its import line alone:
353	# `import ../torch/fast.bend as F` and F.Linear / F.LinearFirst / F.LinearNB, no renames.
354	
355	def Linear(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
356	  LinearF(n, i, o, par)
357	
358	def LinearFirst(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<T.Lin<i, o>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
359	  LinearFirstF(n, i, o, par)
360	
361	def LinearNB(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> T.Layer<TA.Mat<o, i>, TA.Mat<n, i>, TA.Mat<n, o>, TA.Mat<n, i>>:
362	  LinearNBF(n, i, o, par)