tensor.bend

bendtorch@1.2.0.0 · 694 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.14.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      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
85def 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
89def 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
93def Bool.check(b: Bool) -> Maybe<&1, {True{} == b : Bool}>:
94  match b:
95    case True{}:
96      Some{{==}}
97    case False{}:
98      None{}
99
100def 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)
108def 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)
122type Ax is Data:
123  Ax{n: U32, st: U32}
124
125type 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
131def 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
135def 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
143def 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
146def 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)
154def 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
165type P2 is Type:
166  P2{a: Array<F32>, b: Array<F32>}
167
168type 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
172def 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
180def 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)]
188def 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
195def new(+n: Nat) -> Array<F32>:
196  Array.new(F32, TA.cap_depth(n), 0.0)
197
198# =====================================================================================================
199# Operations
200# =====================================================================================================
201
202def zeros(+s: List<&2, Nat>) -> Tensor<s>:
203  Tensor{new(numel(s))}
204
205def 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
210def 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
216type 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
219def 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
224def f32_add(+a: F32, +b: F32) -> F32:
225  (a + b : F32)
226
227def f32_sub(+a: F32, +b: F32) -> F32:
228  (a - b : F32)
229
230def f32_mul(+a: F32, +b: F32) -> F32:
231  (a * b : F32)
232
233def 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
237def 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
242def 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
246type TP<-a: List<&2, Nat>, -b: List<&2, Nat>> is Type:
247  TP{x: Tensor<a>, y: Tensor<b>}
248
249def 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
255def 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
263def 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
272def 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
276def 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
282def 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)
292def 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
301def 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
306def 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
313def 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)
317def 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
325def 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
332def 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
337def 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)
346def 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
354type 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
358type 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
361def 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
367def 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)
373def d_id(+d: F32, +p: F32, +q: F32) -> F32:
374  d
375
376def d_neg(+d: F32, +p: F32, +q: F32) -> F32:
377  (0.0 - d : F32)
378
379def d_mul_a(+d: F32, +p: F32, +q: F32) -> F32:
380  (d * q : F32)
381
382def d_mul_b(+d: F32, +p: F32, +q: F32) -> F32:
383  (d * p : F32)
384
385def 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
389def d_div_b(+d: F32, +p: F32, +q: F32) -> F32:
390  (0.0 - (d * ((p / q : F32) / q : F32) : F32) : F32)
391
392def 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
395def 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
398def 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
406def 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
411def 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
417def 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
421def 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
425def 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)
429def 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
434def 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
439def mean_all(+s: List<&2, Nat>, x: Tensor<s>) -> TP<s, []>:
440  tp_div(s, [], n_f32(numel(s)), sum_all(s, x))
441
442def 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
449def 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
455def 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
458def 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
466def 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
473def 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
481def 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
489def 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
493def 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
500def perm_map(+s: List<&2, Nat>, +p: List<&2, Nat>) -> IMap:
501  Strided{0, perm_axes(p, s, Nil{})}
502
503def 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
509def 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
515def 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
522def 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
530def 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
533def 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
537def 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
543def 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
549def 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
559def 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
570def 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
573def 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
577def cat_map_a(+a: List<&2, Nat>, +b: List<&2, Nat>, +d: Nat) -> IMap:
578  window(a, cat_shape(a, b, d), 0)
579
580def 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
583def 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
588def 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
594def 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
600type 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
603def 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
608def 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
613def 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
623type Geo is Data:
624  Geo{n: U32, m: U32, k: Nat, xb: U32, xi: U32, xt: U32, yb: U32, yt: U32, yj: U32}
625
626type Dot is Type:
627  Dot{x: Array<F32>, y: Array<F32>, v: F32}
628
629def 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
635def 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
642def 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
647def 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
652def geo_k(+g: Geo) -> Nat:
653  match g:
654    case Geo{n, m, +k, xb, xi, xt, yb, yt, yj}:
655      k
656
657def geo_xt(+g: Geo) -> U32:
658  match g:
659    case Geo{n, m, k, xb, xi, +xt, yb, yt, yj}:
660      xt
661
662def 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
668def 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
676def u(+n: Nat) -> U32:
677  U32.from_nat(n)
678
679def 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
685def 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
691def 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)))))