tensor.bend

bendtorch@1.3.0.0 · 701 lines · 120 definitions

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
17import Base
18import bend-ml-tensor-array@0.1.18.0/main.bend as TA
19import bend-ml-nat-lemmas@0.1.2.0/main.bend as NL
20
21# =====================================================================================================
22# Shapes
23# =====================================================================================================
24
25type Tensor<-s: List<&2, Nat>> is Type:
26  Tensor{data: Array<F32>}
27
28# the number of elements of a tensor of shape s
29def 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)
33def 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)
43def 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
53def 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
61def 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
67def 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
71def 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
75def 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
86def 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
90def 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
95def Bool.check(b: Bool) -> Maybe<&1, {True{} == b : Bool}>:
96  match b:
97    case True{}:
98      Some{{==}}
99    case False{}:
100      None{}
101
102def 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)
110def 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)
124type Ax is Data:
125  Ax{n: U32, st: U32}
126
127type 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
133def 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
137def 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
145def 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
148def 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)
156def 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
167type P2 is Type:
168  P2{a: Array<F32>, b: Array<F32>}
169
170type 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
174def 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
182def 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)]
190def 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
197def new(+n: Nat) -> Array<F32>:
198  Array.new(F32, TA.cap_depth(n), 0.0)
199
200# =====================================================================================================
201# Operations
202# =====================================================================================================
203
204def zeros(+s: List<&2, Nat>) -> Tensor<s>:
205  Tensor{new(numel(s))}
206
207def 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
212def 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
218type 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
221def 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
226def f32_add(+a: F32, +b: F32) -> F32:
227  (a + b : F32)
228
229def f32_sub(+a: F32, +b: F32) -> F32:
230  (a - b : F32)
231
232def f32_mul(+a: F32, +b: F32) -> F32:
233  (a * b : F32)
234
235def 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
239def 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
244def 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
248type TP<-a: List<&2, Nat>, -b: List<&2, Nat>> is Type:
249  TP{x: Tensor<a>, y: Tensor<b>}
250
251def 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
257def 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
265def 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
274def 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
278def 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
284def 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)
294def 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
303def 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
308def 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
315def 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)
319def 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
327def 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
334def 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
339def 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)
348def 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
356type 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
360type 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
363def 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
369def 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)
375def d_id(+d: F32, +p: F32, +q: F32) -> F32:
376  d
377
378def d_neg(+d: F32, +p: F32, +q: F32) -> F32:
379  (0.0 - d : F32)
380
381def d_mul_a(+d: F32, +p: F32, +q: F32) -> F32:
382  (d * q : F32)
383
384def d_mul_b(+d: F32, +p: F32, +q: F32) -> F32:
385  (d * p : F32)
386
387def 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
391def d_div_b(+d: F32, +p: F32, +q: F32) -> F32:
392  (0.0 - (d * ((p / q : F32) / q : F32) : F32) : F32)
393
394def 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
397def 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
400def 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
408def 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
413def 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
419def 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
423def 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
427def 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)
431def 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
436def 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
441def mean_all(+s: List<&2, Nat>, x: Tensor<s>) -> TP<s, []>:
442  tp_div(s, [], n_f32(numel(s)), sum_all(s, x))
443
444def 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
451def 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
457def 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
460def 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
468def 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
475def 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
483def 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
492def 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
497def 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
504def perm_map(+s: List<&2, Nat>, +p: List<&2, Nat>) -> IMap:
505  Strided{0, perm_axes(p, s, Nil{})}
506
507def 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
513def 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
519def 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
526def 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
534def 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
537def 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
541def 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
547def 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
553def 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
564def 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
576def 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
580def 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
584def cat_map_a(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat) -> IMap:
585  window(a, cat_shape(a, b, d), 0)
586
587def 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
590def 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
595def 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
601def 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
607type 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
610def 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
615def 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
620def 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
630type Geo is Data:
631  Geo{n: U32, m: U32, k: Nat, xb: U32, xi: U32, xt: U32, yb: U32, yt: U32, yj: U32}
632
633type Dot is Type:
634  Dot{x: Array<F32>, y: Array<F32>, v: F32}
635
636def 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
642def 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
649def 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
654def 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
659def geo_k(+g: Geo) -> Nat:
660  match g:
661    case Geo{n, m, +k, xb, xi, xt, yb, yt, yj}:
662      k
663
664def geo_xt(+g: Geo) -> U32:
665  match g:
666    case Geo{n, m, k, xb, xi, +xt, yb, yt, yj}:
667      xt
668
669def 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
675def 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
683def u(+n: Nat) -> U32:
684  U32.from_nat(n)
685
686def 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
692def 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
698def 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)))))