1	# BendTorch tensors: Tensor<s>, an n-dimensional F32 array whose shape s (a list of sizes) is in its type.
2	#
3	#   import ../torch/tensor.bend as TT
4	#
5	# A Tensor<s> holds product(s) numbers in one flat Array<F32>, row-major (the last axis varies fastest), as bend-ml's
6	# Mat<r, c> does for two axes: Tensor<[r, c]> and Mat<r, c> are the same array under two types, and `of_mat`/`to_mat`
7	# convert without a copy. Every operation hands back what it reads, as the Mat operations do.
8	#
9	# Each operation is one of two kernels over an index map (for each position of one tensor, the position of another):
10	#   gather       out[j] = f(x[m(j)])           (and its two-input form, for broadcasting)
11	#   scatter-add  g[m(j)] = g[m(j)] + x[j]      (reductions, and the gradient of a broadcast)
12	# An index map is a list of axes, each with its size and the stride of the tensor being read (0 on a broadcast axis),
13	# or an axis map (the position with one axis taken out). The shape conditions of an operation are Bool functions of the
14	# shapes (`bc_ok` for broadcasting); a caller passes the proof that the condition is True, which the checker computes for
15	# literal shapes and a program checks once where shapes read at run time enter (`Bool.check`).
16	
17	import Base
18	import bend-ml-tensor-array@0.1.14.0/main.bend as TA
19	import bend-ml-nat-lemmas@0.1.2.0/main.bend as NL
20	
21	# =====================================================================================================
22	# Shapes
23	# =====================================================================================================
24	
25	type Tensor<-s: List<&2, Nat>> is Type:
26	  Tensor{data: Array<F32>}
27	
28	# the number of elements of a tensor of shape s
29	def numel(+s: List<&2, Nat>) -> Nat:
30	  NL.product(s)
31	
32	# size of axis k (1 past the last axis, so a missing axis counts as a size-1 one)
33	def at1(s: List<&2, Nat>, k: Nat) -> Nat:
34	  match s k:
35	    case Nil{} _:
36	      1n
37	    case Con{h, t} 0n:
38	      h
39	    case Con{h, t} 1n+p:
40	      at1(t, p)
41	
42	# s without axis k (s itself past the last axis)
43	def remove(s: List<&2, Nat>, k: Nat) -> List<&2, Nat>:
44	  match s k:
45	    case Nil{} _:
46	      Nil{}
47	    case Con{h, t} 0n:
48	      t
49	    case Con{h, t} 1n+p:
50	      h <> remove(t, p)
51	
52	# n copies of 1, then s: s seen with n more leading axes
53	def ones_then(n: Nat, s: List<&2, Nat>) -> List<&2, Nat>:
54	  match n:
55	    case 0n:
56	      s
57	    case 1n+p:
58	      1n <> ones_then(p, s)
59	
60	# s aligned to rank r by leading size-1 axes, as broadcasting aligns shapes from the right
61	def pad(+s: List<&2, Nat>, +r: Nat) -> List<&2, Nat>:
62	  ones_then(Nat.sub(r, List.length(&2, Nat, s)), s)
63	
64	# ---- broadcasting (PyTorch's rule): shapes aligned from the right; on each axis the sizes are equal or one is 1 ----
65	
66	# one axis: z is x, and y is x or 1
67	def bc_pair(+x: Nat, +y: Nat, +z: Nat) -> Bool:
68	  Bool.and(Nat.is_eq(z, x), Bool.or(Nat.is_eq(y, x), Nat.is_eq(y, 1n)))
69	
70	# one axis, either way round
71	def bc_dim(+x: Nat, +y: Nat, +z: Nat) -> Bool:
72	  Bool.or(bc_pair(x, y, z), bc_pair(y, x, z))
73	
74	# three shapes of one rank, axis by axis
75	def bc_eq(a: List<&2, Nat>, b: List<&2, Nat>, o: List<&2, Nat>) -> Bool:
76	  match a b o:
77	    case Nil{} Nil{} Nil{}:
78	      True{}
79	    case Con{+x, ra} Con{+y, rb} Con{+z, ro}:
80	      Bool.and(bc_dim(x, y, z), bc_eq(ra, rb, ro))
81	    case _ _ _:
82	      False{}
83	
84	# the ranks: o's is the larger of a's and b's
85	def bc_rank(+la: Nat, +lb: Nat, +lo: Nat) -> Bool:
86	  Bool.and(Bool.and(Nat.is_le(la, lo), Nat.is_le(lb, lo)), Bool.or(Nat.is_eq(la, lo), Nat.is_eq(lb, lo)))
87	
88	# a and b broadcast to o
89	def bc_ok(+a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>) -> Bool:
90	  Bool.and(bc_rank(List.length(&2, Nat, a), List.length(&2, Nat, b), List.length(&2, Nat, o)), bc_eq(pad(a, List.length(&2, Nat, o)), pad(b, List.length(&2, Nat, o)), o))
91	
92	# a proof that b is True, when it is: where shapes read at run time enter, a program checks a condition once
93	def Bool.check(b: Bool) -> Maybe<&1, {True{} == b : Bool}>:
94	  match b:
95	    case True{}:
96	      Some{{==}}
97	    case False{}:
98	      None{}
99	
100	def nat_eq_succ(-a: Nat, -b: Nat, m: Maybe<&1, {a == b : Nat}>) -> Maybe<&1, {1n+a == 1n+b : Nat}>:
101	  match m:
102	    case Some{e}:
103	      Some{Equal.cong(Nat, Nat, x => 1n+x, a, b, e)}
104	    case None{}:
105	      None{}
106	
107	# a proof that a == b, when it holds (reshape's count, read at run time)
108	def Nat.check_eq(a: Nat, b: Nat) -> Maybe<&1, {a == b : Nat}>:
109	  match a b:
110	    case 0n 0n:
111	      Some{{==}}
112	    case 1n+p 1n+q:
113	      nat_eq_succ(p, q, Nat.check_eq(p, q))
114	    case _ _:
115	      None{}
116	
117	# =====================================================================================================
118	# Index maps
119	# =====================================================================================================
120	
121	# one axis of an index map: its size, and the stride of the tensor read along it (0 where it is broadcast)
122	type Ax is Data:
123	  Ax{n: U32, st: U32}
124	
125	type IMap is Data:
126	  # base plus, for every axis (the last one first), its coordinate times its stride
127	  Strided{base: U32, axes: List<&2, Ax>}
128	  # the position with one axis taken out: inner is the product of the sizes after it, size its own size
129	  Drop{inner: U32, size: U32}
130	
131	def stride_of(+d: Nat, +st: U32) -> U32:
132	  Bool.pick(U32, Nat.is_eq(d, 1n), 0, st)
133	
134	# the axes of reading a tensor of shape a (aligned to o, same rank) at every position of o; the last axis first
135	def axes(ra: List<&2, Nat>, ro: List<&2, Nat>, +st: U32, acc: List<&2, Ax>) -> List<&2, Ax>:
136	  match ra ro:
137	    case Con{+d, ta} Con{+n, to}:
138	      axes(ta, to, (st * U32.from_nat(d) : U32), Ax{U32.from_nat(n), stride_of(d, st)} <> acc)
139	    case _ _:
140	      acc
141	
142	# reading a (any rank up to o's) along o
143	def bmap(+a: List<&2, Nat>, +o: List<&2, Nat>) -> IMap:
144	  Strided{0, List.reverse(&2, Ax, axes(List.reverse(&2, Nat, pad(a, List.length(&2, Nat, o))), List.reverse(&2, Nat, o), 1, Nil{}))}
145	
146	def strided(axs: List<&2, Ax>, +j: U32, acc: U32) -> U32:
147	  match axs:
148	    case Nil{}:
149	      acc
150	    case Con{Ax{+n, +st}, r}:
151	      strided(r, U32.div(j, n), (acc + (U32.mod(j, n) * st : U32) : U32))
152	
153	# the position m(j)
154	def imap(+m: IMap, +j: U32) -> U32:
155	  match m:
156	    case Strided{+base, axs}:
157	      strided(axs, j, base)
158	    case Drop{+inner, +size}:
159	      ((U32.div(j, (inner * size : U32)) * inner : U32) + U32.mod(j, inner) : U32)
160	
161	# =====================================================================================================
162	# Kernels
163	# =====================================================================================================
164	
165	type P2 is Type:
166	  P2{a: Array<F32>, b: Array<F32>}
167	
168	type P3 is Type:
169	  P3{a: Array<F32>, b: Array<F32>, c: Array<F32>}
170	
171	# o[j] = f(x[m(j)], y[n(j)]) for k positions from j; rx and ry hold x[m(j)] and y[n(j)], already read
172	def gather2_l(~f: @+a: F32 -> @+b: F32 -> F32, k: Nat, +j: U32, +m: IMap, +n: IMap, rx: Array<F32> & F32, ry: Array<F32> & F32, o: Array<F32>) -> P3:
173	  match k rx ry:
174	    case 0n Tuple{x, u} Tuple{y, v}:
175	      P3{x, y, o}
176	    case 1n+p Tuple{x, +u} Tuple{y, +v}:
177	      gather2_l(~f, p, (j + 1 : U32), m, n, Array.get(F32, x, imap(m, (j + 1 : U32))), Array.get(F32, y, imap(n, (j + 1 : U32))), Array.set(F32, o, j, f(u, v)))
178	
179	# o[j] = x[m(j)] for k positions from j
180	def gather_l(k: Nat, +j: U32, +m: IMap, rx: Array<F32> & F32, o: Array<F32>) -> P2:
181	  match k rx:
182	    case 0n Tuple{x, u}:
183	      P2{x, o}
184	    case 1n+p Tuple{x, +u}:
185	      gather_l(p, (j + 1 : U32), m, Array.get(F32, x, imap(m, (j + 1 : U32))), Array.set(F32, o, j, u))
186	
187	# g[m(j)] = g[m(j)] + x[j] for k positions from j; rx holds x[j], rg holds g[m(j)]
188	def scatter_l(k: Nat, +j: U32, +m: IMap, rx: Array<F32> & F32, rg: Array<F32> & F32) -> P2:
189	  match k rx rg:
190	    case 0n Tuple{x, u} Tuple{g, w}:
191	      P2{x, g}
192	    case 1n+p Tuple{x, +u} Tuple{g, +w}:
193	      scatter_l(p, (j + 1 : U32), m, Array.get(F32, x, (j + 1 : U32)), Array.get(F32, Array.set(F32, g, imap(m, j), (w + u : F32)), imap(m, (j + 1 : U32))))
194	
195	def new(+n: Nat) -> Array<F32>:
196	  Array.new(F32, TA.cap_depth(n), 0.0)
197	
198	# =====================================================================================================
199	# Operations
200	# =====================================================================================================
201	
202	def zeros(+s: List<&2, Nat>) -> Tensor<s>:
203	  Tensor{new(numel(s))}
204	
205	def of_mat(-r: Nat, -c: Nat, m: TA.Mat<r, c>) -> Tensor<[r, c]>:
206	  match m:
207	    case TA.Mat{d}:
208	      Tensor{d}
209	
210	def to_mat(-r: Nat, -c: Nat, t: Tensor<[r, c]>) -> TA.Mat<r, c>:
211	  match t:
212	    case Tensor{d}:
213	      TA.Mat{d}
214	
215	# a binary operation, both operands handed back with the result
216	type T3<-a: List<&2, Nat>, -b: List<&2, Nat>, -o: List<&2, Nat>> is Type:
217	  T3{x: Tensor<a>, y: Tensor<b>, z: Tensor<o>}
218	
219	def t3_of(-a: List<&2, Nat>, -b: List<&2, Nat>, -o: List<&2, Nat>, p: P3) -> T3<a, b, o>:
220	  match p:
221	    case P3{x, y, z}:
222	      T3{Tensor{x}, Tensor{y}, Tensor{z}}
223	
224	def f32_add(+a: F32, +b: F32) -> F32:
225	  (a + b : F32)
226	
227	def f32_sub(+a: F32, +b: F32) -> F32:
228	  (a - b : F32)
229	
230	def f32_mul(+a: F32, +b: F32) -> F32:
231	  (a * b : F32)
232	
233	def f32_div(+a: F32, +b: F32) -> F32:
234	  (a / b : F32)
235	
236	# z = f(x, y) element by element, x and y broadcast to o
237	def map2(~f: @+a: F32 -> @+b: F32 -> F32, +a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>, -pf: {True{} == bc_ok(a, b, o) : Bool}, x: Tensor<a>, y: Tensor<b>) -> T3<a, b, o>:
238	  match x y:
239	    case Tensor{dx} Tensor{dy}:
240	      t3_of(a, b, o, gather2_l(~f, numel(o), 0, bmap(a, o), bmap(b, o), Array.get(F32, dx, imap(bmap(a, o), 0)), Array.get(F32, dy, imap(bmap(b, o), 0)), new(numel(o))))
241	
242	def add(+a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>, -pf: {True{} == bc_ok(a, b, o) : Bool}, x: Tensor<a>, y: Tensor<b>) -> T3<a, b, o>:
243	  map2(~f32_add, a, b, o, pf, x, y)
244	
245	# a tensor and the gradient that comes back with it
246	type TP<-a: List<&2, Nat>, -b: List<&2, Nat>> is Type:
247	  TP{x: Tensor<a>, y: Tensor<b>}
248	
249	def tp_of(-a: List<&2, Nat>, -b: List<&2, Nat>, p: P2) -> TP<a, b>:
250	  match p:
251	    case P2{x, y}:
252	      TP{Tensor{x}, Tensor{y}}
253	
254	# g + (dz summed down to a's shape): the gradient of an operand that was broadcast from a to o; dz comes back
255	def sum_to(+a: List<&2, Nat>, +o: List<&2, Nat>, dz: Tensor<o>, g: Tensor<a>) -> TP<o, a>:
256	  match dz g:
257	    case Tensor{d} Tensor{e}:
258	      tp_of(o, a, scatter_l(numel(o), 0, bmap(a, o), Array.get(F32, d, 0), Array.get(F32, e, imap(bmap(a, o), 0))))
259	
260	# ---- sum over one axis ----
261	
262	# the product of the sizes after axis k
263	def inner(s: List<&2, Nat>, k: Nat) -> Nat:
264	  match s k:
265	    case Nil{} _:
266	      1n
267	    case Con{h, t} 0n:
268	      NL.product(t)
269	    case Con{h, t} 1n+p:
270	      inner(t, p)
271	
272	def drop_map(+s: List<&2, Nat>, +k: Nat) -> IMap:
273	  Drop{U32.from_nat(inner(s, k)), U32.from_nat(at1(s, k))}
274	
275	# x summed over axis k: shape remove(s, k); x comes back
276	def sum_axis(+s: List<&2, Nat>, +k: Nat, x: Tensor<s>) -> TP<s, remove(s, k)>:
277	  match x:
278	    case Tensor{d}:
279	      tp_of(s, remove(s, k), scatter_l(numel(s), 0, drop_map(s, k), Array.get(F32, d, 0), Array.get(F32, new(numel(remove(s, k))), 0)))
280	
281	# the gradient of sum_axis: dy (shape remove(s, k)) copied along axis k; dy comes back
282	def sum_axis_bwd(+s: List<&2, Nat>, +k: Nat, dy: Tensor<remove(s, k)>) -> TP<remove(s, k), s>:
283	  match dy:
284	    case Tensor{d}:
285	      tp_of(remove(s, k), s, gather_l(numel(s), 0, drop_map(s, k), Array.get(F32, d, imap(drop_map(s, k), 0)), new(numel(s))))
286	
287	# =====================================================================================================
288	# More shapes
289	# =====================================================================================================
290	
291	# s with the size of axis k replaced by v (s itself past the last axis)
292	def setd(s: List<&2, Nat>, k: Nat, +v: Nat) -> List<&2, Nat>:
293	  match s k:
294	    case Nil{} _:
295	      Nil{}
296	    case Con{h, t} 0n:
297	      v <> t
298	    case Con{h, t} 1n+p:
299	      h <> setd(t, p, v)
300	
301	def rank(+s: List<&2, Nat>) -> Nat:
302	  List.length(&2, Nat, s)
303	
304	# the index map of reading a tensor of strides st (given by its shape st) at every position of a tensor of shape o, from
305	# a base: axis e has o's size and st's stride (inner(st, e)); the last axis first
306	def win_axes(n: Nat, +o: List<&2, Nat>, +st: List<&2, Nat>, +e: Nat, acc: List<&2, Ax>) -> List<&2, Ax>:
307	  match n:
308	    case 0n:
309	      acc
310	    case 1n+p:
311	      win_axes(p, o, st, 1n+e, Ax{U32.from_nat(at1(o, e)), U32.from_nat(inner(st, e))} <> acc)
312	
313	def window(+o: List<&2, Nat>, +st: List<&2, Nat>, +base: U32) -> IMap:
314	  Strided{base, win_axes(rank(o), o, st, 0n, Nil{})}
315	
316	# g[m(j)] = x[j] for k positions from j (a copy into a window of g; no addition, so -0 stays -0)
317	def place_l(k: Nat, +j: U32, +m: IMap, rx: Array<F32> & F32, g: Array<F32>) -> P2:
318	  match k rx:
319	    case 0n Tuple{x, u}:
320	      P2{x, g}
321	    case 1n+p Tuple{x, +u}:
322	      place_l(p, (j + 1 : U32), m, Array.get(F32, x, (j + 1 : U32)), Array.set(F32, g, imap(m, j), u))
323	
324	# a[j] = f(a[j]) for k positions from j, with a constant c
325	def ewc_l(~f: @+a: F32 -> @+c: F32 -> F32, k: Nat, +j: U32, +c: F32, ra: Array<F32> & F32) -> Array<F32>:
326	  match k ra:
327	    case 0n Tuple{a, v}:
328	      a
329	    case 1n+p Tuple{a, +v}:
330	      ewc_l(~f, p, (j + 1 : U32), c, Array.get(F32, Array.set(F32, a, j, f(v, c)), (j + 1 : U32)))
331	
332	def ew_div(+s: List<&2, Nat>, +c: F32, t: Tensor<s>) -> Tensor<s>:
333	  match t:
334	    case Tensor{d}:
335	      Tensor{ewc_l(~f32_div, numel(s), 0, c, Array.get(F32, d, 0))}
336	
337	def n_f32(+n: Nat) -> F32:
338	  U32.to_f32(U32.from_nat(n))
339	
340	# =====================================================================================================
341	# Binary operations with broadcasting: gradients
342	# =====================================================================================================
343	
344	# g[mg(j)] = g[mg(j)] + f(dz[j], x[ma(j)], y[mb(j)]) for k positions from j: the gradient of one operand of map2,
345	# with both operands at hand (mul and div need them)
346	def grad2_l(~f: @+d: F32 -> @+p: F32 -> @+q: F32 -> F32, k: Nat, +j: U32, +mg: IMap, +ma: IMap, +mb: IMap, rd: Array<F32> & F32, rx: Array<F32> & F32, ry: Array<F32> & F32, rg: Array<F32> & F32) -> P4:
347	  match k rd rx ry rg:
348	    case 0n Tuple{d, u} Tuple{x, p} Tuple{y, q} Tuple{g, w}:
349	      P4{d, x, y, g}
350	    case 1n+n Tuple{d, +u} Tuple{x, +p} Tuple{y, +q} Tuple{g, +w}:
351	      grad2_l(~f, n, (j + 1 : U32), mg, ma, mb, Array.get(F32, d, (j + 1 : U32)), Array.get(F32, x, imap(ma, (j + 1 : U32))), Array.get(F32, y, imap(mb, (j + 1 : U32))),
352	        Array.get(F32, Array.set(F32, g, imap(mg, j), (w + f(u, p, q) : F32)), imap(mg, (j + 1 : U32))))
353	
354	type P4 is Type:
355	  P4{a: Array<F32>, b: Array<F32>, c: Array<F32>, d: Array<F32>}
356	
357	# dz, x, y and the gradient, handed back
358	type G2<-a: List<&2, Nat>, -b: List<&2, Nat>, -o: List<&2, Nat>, -g: List<&2, Nat>> is Type:
359	  G2{dz: Tensor<o>, x: Tensor<a>, y: Tensor<b>, g: Tensor<g>}
360	
361	def g2_of(-a: List<&2, Nat>, -b: List<&2, Nat>, -o: List<&2, Nat>, -g: List<&2, Nat>, p: P4) -> G2<a, b, o, g>:
362	  match p:
363	    case P4{d, x, y, w}:
364	      G2{Tensor{d}, Tensor{x}, Tensor{y}, Tensor{w}}
365	
366	# g + the gradient, through f, of the operand of shape gs (a or b) of z = op(x, y): f(dz, x, y) at each position of o
367	def grad2(~f: @+d: F32 -> @+p: F32 -> @+q: F32 -> F32, +a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>, +gs: List<&2, Nat>, dz: Tensor<o>, x: Tensor<a>, y: Tensor<b>, g: Tensor<gs>) -> G2<a, b, o, gs>:
368	  match dz x y g:
369	    case Tensor{dd} Tensor{dx} Tensor{dy} Tensor{dg}:
370	      g2_of(a, b, o, gs, grad2_l(~f, numel(o), 0, bmap(gs, o), bmap(a, o), bmap(b, o), Array.get(F32, dd, 0), Array.get(F32, dx, imap(bmap(a, o), 0)), Array.get(F32, dy, imap(bmap(b, o), 0)), Array.get(F32, dg, imap(bmap(gs, o), 0))))
371	
372	# the partial derivatives, dz times d op / d operand (PyTorch's formulas and order of operations)
373	def d_id(+d: F32, +p: F32, +q: F32) -> F32:
374	  d
375	
376	def d_neg(+d: F32, +p: F32, +q: F32) -> F32:
377	  (0.0 - d : F32)
378	
379	def d_mul_a(+d: F32, +p: F32, +q: F32) -> F32:
380	  (d * q : F32)
381	
382	def d_mul_b(+d: F32, +p: F32, +q: F32) -> F32:
383	  (d * p : F32)
384	
385	def d_div_a(+d: F32, +p: F32, +q: F32) -> F32:
386	  (d / q : F32)
387	
388	# -dz * ((x / y) / y), as PyTorch's div backward for the divisor
389	def d_div_b(+d: F32, +p: F32, +q: F32) -> F32:
390	  (0.0 - (d * ((p / q : F32) / q : F32) : F32) : F32)
391	
392	def sub(+a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>, -pf: {True{} == bc_ok(a, b, o) : Bool}, x: Tensor<a>, y: Tensor<b>) -> T3<a, b, o>:
393	  map2(~f32_sub, a, b, o, pf, x, y)
394	
395	def mul(+a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>, -pf: {True{} == bc_ok(a, b, o) : Bool}, x: Tensor<a>, y: Tensor<b>) -> T3<a, b, o>:
396	  map2(~f32_mul, a, b, o, pf, x, y)
397	
398	def div(+a: List<&2, Nat>, +b: List<&2, Nat>, +o: List<&2, Nat>, -pf: {True{} == bc_ok(a, b, o) : Bool}, x: Tensor<a>, y: Tensor<b>) -> T3<a, b, o>:
399	  map2(~f32_div, a, b, o, pf, x, y)
400	
401	# =====================================================================================================
402	# Reductions
403	# =====================================================================================================
404	
405	# the sum over axis k with the axis kept (size 1): the same numbers as sum_axis, in the same layout
406	def sum_axis_keep(+s: List<&2, Nat>, +k: Nat, x: Tensor<s>) -> TP<s, setd(s, k, 1n)>:
407	  match x:
408	    case Tensor{d}:
409	      tp_of(s, setd(s, k, 1n), scatter_l(numel(s), 0, drop_map(s, k), Array.get(F32, d, 0), Array.get(F32, new(numel(remove(s, k))), 0)))
410	
411	def tp_div(-s: List<&2, Nat>, +u: List<&2, Nat>, +c: F32, r: TP<s, u>) -> TP<s, u>:
412	  match r:
413	    case TP{x, y}:
414	      TP{x, ew_div(u, c, y)}
415	
416	# the mean over axis k: the sum divided by the axis' size
417	def mean_axis(+s: List<&2, Nat>, +k: Nat, x: Tensor<s>) -> TP<s, remove(s, k)>:
418	  tp_div(s, remove(s, k), n_f32(at1(s, k)), sum_axis(s, k, x))
419	
420	# the gradient of mean_axis: dy divided by the axis' size, copied along it
421	def mean_axis_bwd(+s: List<&2, Nat>, +k: Nat, dy: Tensor<remove(s, k)>) -> TP<remove(s, k), s>:
422	  tp_div(remove(s, k), s, n_f32(at1(s, k)), sum_axis_bwd(s, k, dy))
423	
424	# every position to position 0: the whole-tensor reductions
425	def all_map(+s: List<&2, Nat>) -> IMap:
426	  Drop{1, U32.from_nat(numel(s))}
427	
428	# the sum of every element: a tensor of shape [] (one number)
429	def sum_all(+s: List<&2, Nat>, x: Tensor<s>) -> TP<s, []>:
430	  match x:
431	    case Tensor{d}:
432	      tp_of(s, [], scatter_l(numel(s), 0, all_map(s), Array.get(F32, d, 0), Array.get(F32, new(1n), 0)))
433	
434	def sum_all_bwd(+s: List<&2, Nat>, dy: Tensor<[]>) -> TP<[], s>:
435	  match dy:
436	    case Tensor{d}:
437	      tp_of([], s, gather_l(numel(s), 0, all_map(s), Array.get(F32, d, 0), new(numel(s))))
438	
439	def mean_all(+s: List<&2, Nat>, x: Tensor<s>) -> TP<s, []>:
440	  tp_div(s, [], n_f32(numel(s)), sum_all(s, x))
441	
442	def mean_all_bwd(+s: List<&2, Nat>, dy: Tensor<[]>) -> TP<[], s>:
443	  tp_div([], s, n_f32(numel(s)), sum_all_bwd(s, dy))
444	
445	# =====================================================================================================
446	# Reshaping: the same numbers under another shape, with a proof that the count does not change
447	# =====================================================================================================
448	
449	def reshape(+s: List<&2, Nat>, +t: List<&2, Nat>, -e: {numel(s) == numel(t) : Nat}, x: Tensor<s>) -> Tensor<t>:
450	  match x:
451	    case Tensor{d}:
452	      Tensor{d}
453	
454	# n * 1 == n, so a shape and its flattening hold the same count
455	def flat_numel(+s: List<&2, Nat>) -> {numel(s) == numel([numel(s)]) : Nat}:
456	  Equal.sym(Nat, Nat.mul(numel(s), 1n), numel(s), NL.mul_one_r(numel(s)))
457	
458	def flatten(+s: List<&2, Nat>, x: Tensor<s>) -> Tensor<[numel(s)]>:
459	  reshape(s, [numel(s)], flat_numel(s), x)
460	
461	# =====================================================================================================
462	# Permutations: a copy with the axes reordered
463	# =====================================================================================================
464	
465	# the shape with axes reordered by p: axis d of the result is axis p[d] of s
466	def permuted(+s: List<&2, Nat>, p: List<&2, Nat>) -> List<&2, Nat>:
467	  match p:
468	    case Nil{}:
469	      Nil{}
470	    case Con{k, r}:
471	      at1(s, k) <> permuted(s, r)
472	
473	def count(+k: Nat, p: List<&2, Nat>) -> Nat:
474	  match p:
475	    case Nil{}:
476	      0n
477	    case Con{+h, r}:
478	      Nat.add(Bool.pick(Nat, Nat.is_eq(h, k), 1n, 0n), count(k, r))
479	
480	# every axis 0..n-1 exactly once, from axis i
481	def all_once(left: Nat, +p: List<&2, Nat>, +i: Nat) -> Bool:
482	  match left:
483	    case 0n:
484	      True{}
485	    case 1n+r:
486	      Bool.and(Nat.is_eq(count(i, p), 1n), all_once(r, p, 1n+i))
487	
488	# p is a permutation of the axes of a rank-n tensor
489	def is_perm(+p: List<&2, Nat>, +n: Nat) -> Bool:
490	  Bool.and(Nat.is_eq(List.length(&2, Nat, p), n), all_once(n, p, 0n))
491	
492	# reading s at every position of the permuted shape: axis d has s's size and s's stride of axis p[d]; the last first
493	def perm_axes(p: List<&2, Nat>, +s: List<&2, Nat>, acc: List<&2, Ax>) -> List<&2, Ax>:
494	  match p:
495	    case Nil{}:
496	      acc
497	    case Con{+k, r}:
498	      perm_axes(r, s, Ax{U32.from_nat(at1(s, k)), U32.from_nat(inner(s, k))} <> acc)
499	
500	def perm_map(+s: List<&2, Nat>, +p: List<&2, Nat>) -> IMap:
501	  Strided{0, perm_axes(p, s, Nil{})}
502	
503	def permute(+s: List<&2, Nat>, +p: List<&2, Nat>, -pf: {True{} == is_perm(p, rank(s)) : Bool}, x: Tensor<s>) -> TP<s, permuted(s, p)>:
504	  match x:
505	    case Tensor{d}:
506	      tp_of(s, permuted(s, p), gather_l(numel(s), 0, perm_map(s, p), Array.get(F32, d, imap(perm_map(s, p), 0)), new(numel(s))))
507	
508	# the gradient of permute: dy put back where each number came from
509	def permute_bwd(+s: List<&2, Nat>, +p: List<&2, Nat>, dy: Tensor<permuted(s, p)>) -> TP<permuted(s, p), s>:
510	  match dy:
511	    case Tensor{d}:
512	      tp_of(permuted(s, p), s, place_l(numel(s), 0, perm_map(s, p), Array.get(F32, d, 0), new(numel(s))))
513	
514	# 0..n-1 with i and j swapped
515	def swap_perm(n: Nat, +i: Nat, +a: Nat, +b: Nat) -> List<&2, Nat>:
516	  match n:
517	    case 0n:
518	      Nil{}
519	    case 1n+r:
520	      Bool.pick(Nat, Nat.is_eq(i, a), b, Bool.pick(Nat, Nat.is_eq(i, b), a, i)) <> swap_perm(r, 1n+i, a, b)
521	
522	def transpose_p(+s: List<&2, Nat>, +i: Nat, +j: Nat) -> List<&2, Nat>:
523	  swap_perm(rank(s), 0n, i, j)
524	
525	# =====================================================================================================
526	# Windows: narrow (a range of one axis) and cat (two tensors along one axis)
527	# =====================================================================================================
528	
529	# axis d exists and start + len fits in it
530	def narrow_ok(+s: List<&2, Nat>, +d: Nat, +st: Nat, +l: Nat) -> Bool:
531	  Bool.and(Nat.is_lt(d, rank(s)), Nat.is_le(Nat.add(st, l), at1(s, d)))
532	
533	def narrow_map(+s: List<&2, Nat>, +d: Nat, +st: Nat, +l: Nat) -> IMap:
534	  window(setd(s, d, l), s, U32.from_nat(Nat.mul(st, inner(s, d))))
535	
536	# positions st .. st + l - 1 of axis d
537	def narrow(+s: List<&2, Nat>, +d: Nat, +st: Nat, +l: Nat, -pf: {True{} == narrow_ok(s, d, st, l) : Bool}, x: Tensor<s>) -> TP<s, setd(s, d, l)>:
538	  match x:
539	    case Tensor{dx}:
540	      tp_of(s, setd(s, d, l), gather_l(numel(setd(s, d, l)), 0, narrow_map(s, d, st, l), Array.get(F32, dx, imap(narrow_map(s, d, st, l), 0)), new(numel(setd(s, d, l)))))
541	
542	# the gradient of narrow: dy in its window, zeros elsewhere
543	def narrow_bwd(+s: List<&2, Nat>, +d: Nat, +st: Nat, +l: Nat, dy: Tensor<setd(s, d, l)>) -> TP<setd(s, d, l), s>:
544	  match dy:
545	    case Tensor{dd}:
546	      tp_of(setd(s, d, l), s, place_l(numel(setd(s, d, l)), 0, narrow_map(s, d, st, l), Array.get(F32, dd, 0), new(numel(s))))
547	
548	# the same sizes on every axis
549	def same(a: List<&2, Nat>, b: List<&2, Nat>) -> Bool:
550	  match a b:
551	    case Nil{} Nil{}:
552	      True{}
553	    case Con{+x, ra} Con{+y, rb}:
554	      Bool.and(Nat.is_eq(x, y), same(ra, rb))
555	    case _ _:
556	      False{}
557	
558	# a and b have the same rank, and the same sizes on every axis but d
559	def cat_eq(a: List<&2, Nat>, b: List<&2, Nat>, d: Nat) -> Bool:
560	  match a b d:
561	    case Nil{} Nil{} _:
562	      True{}
563	    case Con{x, ra} Con{y, rb} 0n:
564	      same(ra, rb)
565	    case Con{+x, ra} Con{+y, rb} 1n+p:
566	      Bool.and(Nat.is_eq(x, y), cat_eq(ra, rb, p))
567	    case _ _ _:
568	      False{}
569	
570	def cat_ok(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat) -> Bool:
571	  Bool.and(Nat.is_lt(d, rank(a)), cat_eq(a, b, d))
572	
573	def cat_shape(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat) -> List<&2, Nat>:
574	  setd(a, d, Nat.add(at1(a, d), at1(b, d)))
575	
576	# where a's and b's positions go in the result
577	def cat_map_a(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat) -> IMap:
578	  window(a, cat_shape(a, b, d), 0)
579	
580	def cat_map_b(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat) -> IMap:
581	  window(b, cat_shape(a, b, d), U32.from_nat(Nat.mul(at1(a, d), inner(cat_shape(a, b, d), d))))
582	
583	def t3_cat(-a: List<&2, Nat>, -b: List<&2, Nat>, -o: List<&2, Nat>, ax: Array<F32>, r: P2) -> T3<a, b, o>:
584	  match r:
585	    case P2{bx, o2}:
586	      T3{Tensor{ax}, Tensor{bx}, Tensor{o2}}
587	
588	def cat_b(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat, rb: Array<F32>, pa: P2) -> T3<a, b, cat_shape(a, b, d)>:
589	  match pa:
590	    case P2{ax, o}:
591	      t3_cat(a, b, cat_shape(a, b, d), ax, place_l(numel(b), 0, cat_map_b(a, b, d), Array.get(F32, rb, 0), o))
592	
593	# a then b along axis d
594	def cat(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat, -pf: {True{} == cat_ok(a, b, d) : Bool}, x: Tensor<a>, y: Tensor<b>) -> T3<a, b, cat_shape(a, b, d)>:
595	  match x y:
596	    case Tensor{dx} Tensor{dy}:
597	      cat_b(a, b, d, dy, place_l(numel(a), 0, cat_map_a(a, b, d), Array.get(F32, dx, 0), new(numel(cat_shape(a, b, d)))))
598	
599	# the gradients of cat: dz's window of a, and dz's window of b (dz comes back); returned as TP<a, b> after dz
600	type CatG<-a: List<&2, Nat>, -b: List<&2, Nat>, -o: List<&2, Nat>> is Type:
601	  CatG{dz: Tensor<o>, ga: Tensor<a>, gb: Tensor<b>}
602	
603	def catg_c(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat, ga: Array<F32>, r: P2) -> CatG<a, b, cat_shape(a, b, d)>:
604	  match r:
605	    case P2{dz, gb}:
606	      CatG{Tensor{dz}, Tensor{ga}, Tensor{gb}}
607	
608	def catg_b(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat, r: P2) -> CatG<a, b, cat_shape(a, b, d)>:
609	  match r:
610	    case P2{dz, ga}:
611	      catg_c(a, b, d, ga, gather_l(numel(b), 0, cat_map_b(a, b, d), Array.get(F32, dz, imap(cat_map_b(a, b, d), 0)), new(numel(b))))
612	
613	def cat_bwd(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat, dz: Tensor<cat_shape(a, b, d)>) -> CatG<a, b, cat_shape(a, b, d)>:
614	  match dz:
615	    case Tensor{dd}:
616	      catg_b(a, b, d, gather_l(numel(a), 0, cat_map_a(a, b, d), Array.get(F32, dd, imap(cat_map_a(a, b, d), 0)), new(numel(a))))
617	
618	# =====================================================================================================
619	# Batched matrix product: [bs, n, k] · [bs, k, m] = [bs, n, m]
620	# =====================================================================================================
621	
622	# where the operands of out[b, i, j] = sum over t of X[b·xb + i·xi + t·xt] · Y[b·yb + t·yt + j·yj] are
623	type Geo is Data:
624	  Geo{n: U32, m: U32, k: Nat, xb: U32, xi: U32, xt: U32, yb: U32, yt: U32, yj: U32}
625	
626	type Dot is Type:
627	  Dot{x: Array<F32>, y: Array<F32>, v: F32}
628	
629	def Dot.then(-R: Type, d: Dot, k: Array<F32> -> Array<F32> -> F32 -> R) -> R:
630	  match d:
631	    case Dot{x, y, v}:
632	      k(x, y, v)
633	
634	# acc + the k products from x[xo] and y[yo], stepping xt and yt; rx and ry hold the first pair, already read
635	def dot_l(k: Nat, +xo: U32, +xt: U32, +yo: U32, +yt: U32, rx: Array<F32> & F32, ry: Array<F32> & F32, acc: F32) -> Dot:
636	  match k rx ry:
637	    case 0n Tuple{x, u} Tuple{y, v}:
638	      Dot{x, y, acc}
639	    case 1n+p Tuple{x, +u} Tuple{y, +v}:
640	      dot_l(p, (xo + xt : U32), xt, (yo + yt : U32), yt, Array.get(F32, x, (xo + xt : U32)), Array.get(F32, y, (yo + yt : U32)), (acc + (u * v : F32) : F32))
641	
642	def x_off(+g: Geo, +j: U32) -> U32:
643	  match g:
644	    case Geo{+n, +m, k, +xb, +xi, xt, yb, yt, yj}:
645	      ((U32.div(j, (n * m : U32)) * xb : U32) + (U32.mod(U32.div(j, m), n) * xi : U32) : U32)
646	
647	def y_off(+g: Geo, +j: U32) -> U32:
648	  match g:
649	    case Geo{+n, +m, k, xb, xi, xt, +yb, yt, +yj}:
650	      ((U32.div(j, (n * m : U32)) * yb : U32) + (U32.mod(j, m) * yj : U32) : U32)
651	
652	def geo_k(+g: Geo) -> Nat:
653	  match g:
654	    case Geo{n, m, +k, xb, xi, xt, yb, yt, yj}:
655	      k
656	
657	def geo_xt(+g: Geo) -> U32:
658	  match g:
659	    case Geo{n, m, k, xb, xi, +xt, yb, yt, yj}:
660	      xt
661	
662	def geo_yt(+g: Geo) -> U32:
663	  match g:
664	    case Geo{n, m, k, xb, xi, xt, yb, +yt, yj}:
665	      yt
666	
667	# out[j] for cnt positions from j
668	def bmm_l(cnt: Nat, +j: U32, +g: Geo, x: Array<F32>, y: Array<F32>, o: Array<F32>) -> P3:
669	  match cnt:
670	    case 0n:
671	      P3{x, y, o}
672	    case 1n+p:
673	      Dot.then(P3, dot_l(geo_k(g), x_off(g, j), geo_xt(g), y_off(g, j), geo_yt(g), Array.get(F32, x, x_off(g, j)), Array.get(F32, y, y_off(g, j)), 0.0),
674	        x2 => y2 => v => bmm_l(p, (j + 1 : U32), g, x2, y2, Array.set(F32, o, j, v)))
675	
676	def u(+n: Nat) -> U32:
677	  U32.from_nat(n)
678	
679	def bmm(+bs: Nat, +n: Nat, +k: Nat, +m: Nat, x: Tensor<[bs, n, k]>, y: Tensor<[bs, k, m]>) -> T3<[bs, n, k], [bs, k, m], [bs, n, m]>:
680	  match x y:
681	    case Tensor{dx} Tensor{dy}:
682	      t3_of([bs, n, k], [bs, k, m], [bs, n, m], bmm_l(Nat.mul(bs, Nat.mul(n, m)), 0, Geo{u(n), u(m), k, u(Nat.mul(n, k)), u(k), 1, u(Nat.mul(k, m)), u(m), 1}, dx, dy, new(Nat.mul(bs, Nat.mul(n, m)))))
683	
684	# dX = dZ · Yᵀ per batch: [bs, n, m] · [bs, k, m]ᵀ = [bs, n, k]; dZ and Y come back
685	def bmm_dx(+bs: Nat, +n: Nat, +k: Nat, +m: Nat, dz: Tensor<[bs, n, m]>, y: Tensor<[bs, k, m]>) -> T3<[bs, n, m], [bs, k, m], [bs, n, k]>:
686	  match dz y:
687	    case Tensor{dd} Tensor{dy}:
688	      t3_of([bs, n, m], [bs, k, m], [bs, n, k], bmm_l(Nat.mul(bs, Nat.mul(n, k)), 0, Geo{u(n), u(k), m, u(Nat.mul(n, m)), u(m), 1, u(Nat.mul(k, m)), 1, u(m)}, dd, dy, new(Nat.mul(bs, Nat.mul(n, k)))))
689	
690	# dY = Xᵀ · dZ per batch: [bs, n, k]ᵀ · [bs, n, m] = [bs, k, m]; X and dZ come back
691	def bmm_dy(+bs: Nat, +n: Nat, +k: Nat, +m: Nat, x: Tensor<[bs, n, k]>, dz: Tensor<[bs, n, m]>) -> T3<[bs, n, k], [bs, n, m], [bs, k, m]>:
692	  match x dz:
693	    case Tensor{dx} Tensor{dd}:
694	      t3_of([bs, n, k], [bs, n, m], [bs, k, m], bmm_l(Nat.mul(bs, Nat.mul(k, m)), 0, Geo{u(k), u(m), n, u(Nat.mul(n, k)), 1, u(k), u(Nat.mul(n, m)), u(m), 1}, dx, dd, new(Nat.mul(bs, Nat.mul(k, m)))))