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.0.1.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	# SGD splits two arrays, momentum three, Adam four. No unused parameter-sized buffers.
88	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>:
89	  match o w g st:
90	    case T.SGD{+lr} TA.Mat{wd} TA.Mat{gd} _:
91	      sgd_fin(r, c, TF.par2(~F32, ~F32, ~F32, ~sgd_leaf, lr, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), wd, gd))
92	    case T.Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} T.OneMom{TA.Mat{md}}:
93	      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))
94	    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}}:
95	      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))
96	    case _ w g st: T.MatStep{w, g, st}
97	
98	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>>:
99	  T.mat_step_fin(r, c, opt_mat(r, c, op, p, g, st))
100	
101	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>>:
102	  match p g st:
103	    case T.Lin{w, b} T.Lin{gw, gb} T.Two{sw, sb}:
104	      T.lin_s2(i, o, opt_mat(o, i, op, w, gw, sw), opt_mat(1n, o, op, b, gb, sb))
105	
106	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>>:
107	  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))))
108	
109	# ---- element-wise maps, zips and copies ----
110	
111	def ew1_leaf(~f: @+a: F32 -> F32, +e: Unit, +i0: U32, +cnt: Nat, a: Array<F32>) -> Array<F32>:
112	  T.ew1_l(~f, cnt, i0, Array.get(F32, a, i0))
113	
114	def map(~f: @+a: F32 -> F32, +r: Nat, +c: Nat, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
115	  match m:
116	    case TA.Mat{d}:
117	      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)}
118	
119	
120	def ew2_leaf(~f: @+a: F32 -> @+b: F32 -> F32, +e: Unit, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
121	  match r:
122	    case TF.B2{a, b}:
123	      p2_b2(T.ew2_l(~f, cnt, i0, Array.get(F32, a, i0), Array.get(F32, b, i0)))
124	
125	def zip_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> TA.MPair<r, c>:
126	  match q:
127	    case TF.B2{a, b}:
128	      TA.MPair{TA.Mat{a}, TA.Mat{b}}
129	
130	# a[i] = f(a[i], b[i]); b comes back unchanged
131	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>:
132	  match a b:
133	    case TA.Mat{x} TA.Mat{y}:
134	      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))
135	
136	def a2_b2(q: T.A2) -> TF.B2<F32, F32>:
137	  match q:
138	    case T.A2{a, b}:
139	      TF.B2{a, b}
140	
141	# numbers i0 .. i0+cnt-1 of a into b
142	def copy_leaf(+e: Unit, +i0: U32, +cnt: Nat, r: TF.B2<F32, F32>) -> TF.B2<F32, F32>:
143	  match cnt r:
144	    case 0n r:
145	      r
146	    case 1n+k TF.B2{a, b}:
147	      a2_b2(T.copy_n(k, i0, i0, b, Array.get(F32, a, i0)))
148	
149	def dup_fin(-r: Nat, -c: Nat, q: TF.B2<F32, F32>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
150	  match q:
151	    case TF.B2{a, b}:
152	      T.Two{TA.Mat{a}, TA.Mat{b}}
153	
154	# two copies of a matrix (T.Mat.dup)
155	def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat<r, c>) -> T.Two<TA.Mat<r, c>, TA.Mat<r, c>>:
156	  match m:
157	    case TA.Mat{a}:
158	      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))))
159	
160	def mpair_a(-r: Nat, -c: Nat, q: TA.MPair<r, c>) -> TA.Mat<r, c>:
161	  match q:
162	    case TA.MPair{a, b}:
163	      a
164	
165	# a + b, dropping b (T.Mat.sum2)
166	def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
167	  mpair_a(r, c, zip(~T.f32_add, r, c, a, b))
168	
169	# ---- Linear: y = x·Wᵀ + b ----
170	
171	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>>:
172	  match p:
173	    case T.Lin{w, b}:
174	      T.lin_f2(n, i, o, b, TF.matmul_nt(n, i, o, par, x, w))
175	
176	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>>:
177	  match g r:
178	    case T.Lin{gw, gb} TA.MMul{dy, w, dx}:
179	      T.lin_b3(n, i, o, w, b, dx, gb, TF.matmul_tn_acc(o, n, i, par, dy, x, gw))
180	
181	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>>:
182	  match p:
183	    case T.Lin{w, b}:
184	      lin_b2(n, i, o, par, b, x, g, TF.matmul_precise(n, o, i, par, dy, w))
185	
186	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>>:
187	  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))}
188	
189	# ---- LinearFirst: no dX (the first layer's input needs no gradient) ----
190	
191	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>>:
192	  match p g:
193	    case T.Lin{w, b} T.Lin{gw, gb}:
194	      T.linf_b3(n, i, o, w, b, gb, TF.matmul_tn_acc(o, n, i, par, dy, x, gw))
195	
196	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>>:
197	  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))}
198	
199	# ---- LinearNB: y = x·Wᵀ, no bias (Llama's projections) ----
200	
201	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>>:
202	  match r:
203	    case TA.MMul{dy, w, dx}:
204	      T.lnb_b3(n, i, o, w, dx, TF.matmul_tn_acc(o, n, i, par, dy, x, g))
205	
206	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>>:
207	  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))}
208	
209	# ---- element-wise layers ----
210	
211	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>>:
212	  match m:
213	    case T.Two{x, s}:
214	      T.Fwd{Unit{}, map(~f, n, c, x), s}
215	
216	# T.ActIn: f on a copy, the input saved
217	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>>:
218	  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}
219	
220	def SiLU(+n: Nat, +c: Nat) -> T.Layer<Unit, TA.Mat<n, c>, TA.Mat<n, c>, TA.Mat<n, c>>:
221	  ActIn(~T.silu_v, ~T.silu_d, n, c)
222	
223	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>:
224	  match x:
225	    case T.Two{a, b}:
226	      T.Fwd{Unit{}, Mat.sum2(n, c, a, b), Unit{}}
227	
228	def Add2(+n: Nat, +c: Nat) -> T.Layer<Unit, T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, TA.Mat<n, c>, Unit>:
229	  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}
230	
231	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>>>:
232	  match da db:
233	    case T.Two{a1, a2} T.Two{b1, b2}:
234	      T.mul2_f3(n, c, a2, b2, zip(~T.f32_mul, n, c, a1, b1))
235	
236	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>>>:
237	  match x:
238	    case T.Two{a, b}:
239	      mul2_f2(n, c, Mat.dup(n, c, a), Mat.dup(n, c, b))
240	
241	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>>>:
242	  match s d:
243	    case T.Two{a, b} T.Two{d1, d2}:
244	      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{}}
245	
246	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>>>:
247	  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}
248	
249	# T.Residual with these copies and sums
250	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>>:
251	  T.Seq(~T.Two<Unit, P>, ~Unit, ~TA.Mat<n, c>, ~T.Two<TA.Mat<n, c>, TA.Mat<n, c>>, ~TA.Mat<n, c>, ~T.Two<Unit, S>, ~Unit, ~T.Fork(~Unit, ~P, ~TA.Mat<n, c>, ~TA.Mat<n, c>, ~TA.Mat<n, c>, ~Unit, ~S, ~(m => Mat.dup(n, c, m)), ~(a => b => Mat.sum2(n, c, a, b)), ~T.Id(~TA.Mat<n, c>), ~f), ~Add2(n, c))
252	
253	# main.bend's RMSNorm, RoPE, Attention, MSE, Embedding and CrossEntropy under their own names, so that a program that
254	# takes its layers from here can take device versions of the same names from another file by its import line alone
255	# (they run as in main.bend)
256	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>>:
257	  T.RMSNorm(n, d, eps)
258	
259	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>:
260	  T.RoPE(n, w, t, h, half, base, e)
261	
262	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>:
263	  T.Attention(n, w, kvw, b, t, kv, g, hd, sc, e1, e2, e3)
264	
265	def MSE(+n: Nat, +c: Nat, y: TA.Mat<n, c>, t: TA.Mat<n, c>) -> T.LossOut<TA.Mat<n, c>>:
266	  T.MSE(n, c, y, t)
267	
268	def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
269	  T.Embedding(n, v, d)
270	
271	def CrossEntropy(+n: Nat, +c: Nat, z: TA.Mat<n, c>, +labels: List<&2, Nat>) -> T.LossOut<TA.Mat<n, c>>:
272	  T.CrossEntropy(n, c, z, labels)