tensor.bend
bendtorch@1.2.0.0 · 694 lines · 0 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)))))