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