src/shape/core.bend
bendtorch-core@0.1.0.0 · 927 lines · 137 definitions
1import Base
2
3type Fin.Step<-k: Nat> is Data:
4 FZero{}
5 FNext{previous: Fin(k)}
6
7def Fin(n: Nat) -> Data:
8 match n:
9 case 0n:
10 Empty
11 case 1n+k:
12 Fin.Step<k>
13
14def Fin.value(+n: Nat, index: Fin(n)) -> Nat:
15 match n:
16 case 0n:
17 match index:
18 case 1n+k:
19 match index:
20 case FZero{}:
21 0n
22 case FNext{previous}:
23 1n+Fin.value(k, previous)
24
25def LT(a: Nat, b: Nat) -> Data:
26 match a b:
27 case 0n 0n:
28 Empty
29 case 0n 1n+q:
30 Unit
31 case 1n+p 0n:
32 Empty
33 case 1n+p 1n+q:
34 LT(p, q)
35
36def LE(a: Nat, b: Nat) -> Data:
37 match a b:
38 case 0n q:
39 Unit
40 case 1n+p 0n:
41 Empty
42 case 1n+p 1n+q:
43 LE(p, q)
44
45def Fin.bound(+n: Nat, +index: Fin(n)) -> LT(Fin.value(n, index), n):
46 match n:
47 case 0n:
48 match index:
49 case 1n+k:
50 match index:
51 case FZero{}:
52 Unit{}
53 case FNext{+previous}:
54 Fin.bound(k, previous)
55
56def Fin.left(+a: Nat, -b: Nat, index: Fin(a)) -> Fin(Nat.add(a, b)):
57 match a:
58 case 0n:
59 match index:
60 case 1n+p:
61 match index:
62 case FZero{}:
63 FZero{}
64 case FNext{previous}:
65 FNext{Fin.left(p, b, previous)}
66
67def Fin.shift(a: Nat, -b: Nat, index: Fin(b)) -> Fin(Nat.add(a, b)):
68 match a:
69 case 0n:
70 index
71 case 1n+p:
72 FNext{Fin.shift(p, b, index)}
73
74def Fin.split.step(-a: Nat, -b: Nat, split: Either<Fin(a), Fin(b)>) -> Either<Fin(1n+a), Fin(b)>:
75 match split:
76 case Inl{index}:
77 Inl{FNext{index}}
78 case Inr{index}:
79 Inr{index}
80
81def Fin.split(+a: Nat, +b: Nat, index: Fin(Nat.add(a, b))) -> Either<Fin(a), Fin(b)>:
82 match a:
83 case 0n:
84 Inr{index}
85 case 1n+p:
86 match index:
87 case FZero{}:
88 Inl{FZero{}}
89 case FNext{previous}:
90 Fin.split.step(p, b, Fin.split(p, b, previous))
91
92def Fin.merge(+a: Nat, -b: Nat, split: Either<Fin(a), Fin(b)>) -> Fin(Nat.add(a, b)):
93 match split:
94 case Inl{index}:
95 Fin.left(a, b, index)
96 case Inr{index}:
97 Fin.shift(a, b, index)
98
99def Fin.pack(+a: Nat, +b: Nat, coordinate: Fin(a) & Fin(b)) -> Fin(Nat.mul(a, b)):
100 match a:
101 case 0n:
102 (head, tail) = coordinate
103 match head:
104 case 1n+p:
105 (head, tail) = coordinate
106 match head:
107 case FZero{}:
108 Fin.left(b, Nat.mul(p, b), tail)
109 case FNext{previous}:
110 Fin.shift(b, Nat.mul(p, b), Fin.pack(p, b, (previous, tail)))
111
112def Fin.unpack.lift(-a: Nat, -b: Nat, coordinate: Fin(a) & Fin(b)) -> Fin(1n+a) & Fin(b):
113 (head, tail) = coordinate
114 (FNext{head}, tail)
115
116def Fin.unpack.step(-a: Nat, -b: Nat, rec: Fin(Nat.mul(a, b)) -> Fin(a) & Fin(b), split: Either<Fin(b), Fin(Nat.mul(a, b))>) -> Fin(1n+a) & Fin(b):
117 match split:
118 case Inl{tail}:
119 (FZero{}, tail)
120 case Inr{index}:
121 Fin.unpack.lift(a, b, rec(index))
122
123def Fin.unpack(+a: Nat, +b: Nat, index: Fin(Nat.mul(a, b))) -> Fin(a) & Fin(b):
124 match a:
125 case 0n:
126 match index:
127 case 1n+p:
128 Fin.unpack.step(p, b, index => Fin.unpack(p, b, index), Fin.split(b, Nat.mul(p, b), index))
129
130def numel(+dims: List<&2, Nat>) -> Nat:
131 match dims:
132 case Nil{}:
133 1n
134 case head <> tail:
135 Nat.mul(head, numel(tail))
136
137type Coordinate.Step<-head: Nat, -tail: List<&2, Nat>> is Data:
138 CStep{index: Fin(head), rest: Coordinate(tail)}
139
140def Coordinate(dims: List<&2, Nat>) -> Data:
141 match dims:
142 case Nil{}:
143 Unit
144 case head <> tail:
145 Coordinate.Step<head, tail>
146
147def flatten(+dims: List<&2, Nat>, coordinate: Coordinate(dims)) -> Fin(numel(dims)):
148 match dims:
149 case Nil{}:
150 FZero{}
151 case +head <> +tail:
152 CStep{index, rest} = coordinate
153 Fin.pack(head, numel(tail), (index, flatten(tail, rest)))
154
155def unflatten.step(-head: Nat, -tail: List<&2, Nat>, rec: Fin(numel(tail)) -> Coordinate(tail), coordinate: Fin(head) & Fin(numel(tail))) -> Coordinate(head <> tail):
156 (first, rest) = coordinate
157 CStep{first, rec(rest)}
158
159def unflatten(+dims: List<&2, Nat>, index: Fin(numel(dims))) -> Coordinate(dims):
160 match dims:
161 case Nil{}:
162 Unit{}
163 case +head <> +tail:
164 unflatten.step(head, tail, rest => unflatten(tail, rest), Fin.unpack(head, numel(tail), index))
165
166def row_major(+dims: List<&2, Nat>, coordinate: Coordinate(dims)) -> Nat:
167 match dims:
168 case Nil{}:
169 0n
170 case +head <> +tail:
171 CStep{index, rest} = coordinate
172 Nat.add(Nat.mul(Fin.value(head, index), numel(tail)), row_major(tail, rest))
173
174def indices(+dims: List<&2, Nat>, coordinate: Coordinate(dims)) -> List<&2, Nat>:
175 match dims:
176 case Nil{}:
177 Nil{}
178 case +head <> +tail:
179 CStep{index, rest} = coordinate
180 Fin.value(head, index) <> indices(tail, rest)
181
182def broadcast_index(size: Nat, index: Nat) -> Nat:
183 match size:
184 case 0n:
185 index
186 case 1n+p:
187 match p:
188 case 0n:
189 0n
190 case 1n+q:
191 index
192
193def broadcast_indices(dims: List<&2, Nat>, values: List<&2, Nat>) -> List<&2, Nat>:
194 match dims values:
195 case Nil{} _:
196 Nil{}
197 case dh <> dt Nil{}:
198 Nil{}
199 case dh <> dt vh <> vt:
200 broadcast_index(dh, vh) <> broadcast_indices(dt, vt)
201
202def project_indices(+dims: List<&2, Nat>, +values: List<&2, Nat>) -> List<&2, Nat>:
203 broadcast_indices(dims, List.drop(&2, Nat, values, Nat.sub(List.length(&2, Nat, values), List.length(&2, Nat, dims))))
204
205type Axis<-left: Nat, -right: Nat, -output: Nat> is Data:
206 AxisEqual{left_eq: {output == left : Nat}, right_eq: {output == right : Nat}}
207 AxisLeftOne{one: {1n == left : Nat}, right_eq: {output == right : Nat}}
208 AxisRightOne{left_eq: {output == left : Nat}, one: {1n == right : Nat}}
209
210def Axis.project_left(-left: Nat, -right: Nat, -output: Nat, alignment: Axis<left, right, output>, index: Fin(output)) -> Fin(left):
211 match alignment:
212 case AxisEqual{left_eq, right_eq}:
213 %left_eq : Fin(_)
214 index
215 case AxisLeftOne{one, right_eq}:
216 %one : Fin(_)
217 FZero{}
218 case AxisRightOne{left_eq, one}:
219 %left_eq : Fin(_)
220 index
221
222def Axis.project_right(-left: Nat, -right: Nat, -output: Nat, alignment: Axis<left, right, output>, index: Fin(output)) -> Fin(right):
223 match alignment:
224 case AxisEqual{left_eq, right_eq}:
225 %right_eq : Fin(_)
226 index
227 case AxisLeftOne{one, right_eq}:
228 %right_eq : Fin(_)
229 index
230 case AxisRightOne{left_eq, one}:
231 %one : Fin(_)
232 FZero{}
233
234type Aligned.Step<-lh: Nat, -lt: List<&2, Nat>, -rh: Nat, -rt: List<&2, Nat>, -oh: Nat, -ot: List<&2, Nat>> is Data:
235 AStep{axis: Axis<lh, rh, oh>, tail: Aligned(lt, rt, ot)}
236
237def Aligned(left: List<&2, Nat>, right: List<&2, Nat>, output: List<&2, Nat>) -> Data:
238 match left right output:
239 case Nil{} Nil{} Nil{}:
240 Unit
241 case lh <> lt rh <> rt oh <> ot:
242 Aligned.Step<lh, lt, rh, rt, oh, ot>
243 case Nil{} Nil{} oh <> ot:
244 Empty
245 case Nil{} rh <> rt Nil{}:
246 Empty
247 case lh <> lt Nil{} Nil{}:
248 Empty
249 case Nil{} rh <> rt oh <> ot:
250 Empty
251 case lh <> lt Nil{} oh <> ot:
252 Empty
253 case lh <> lt rh <> rt Nil{}:
254 Empty
255
256def Aligned.project_left(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, alignment: Aligned(left, right, output), coordinate: Coordinate(output)) -> Coordinate(left):
257 match left right output:
258 case Nil{} Nil{} Nil{}:
259 Unit{}
260 case +lh <> +lt +rh <> +rt +oh <> +ot:
261 AStep{axis, tail} = alignment
262 CStep{index, rest} = coordinate
263 CStep{Axis.project_left(lh, rh, oh, axis, index), Aligned.project_left(lt, rt, ot, tail, rest)}
264 case Nil{} Nil{} oh <> ot:
265 match alignment:
266 case Nil{} rh <> rt Nil{}:
267 match alignment:
268 case lh <> lt Nil{} Nil{}:
269 match alignment:
270 case Nil{} rh <> rt oh <> ot:
271 match alignment:
272 case lh <> lt Nil{} oh <> ot:
273 match alignment:
274 case lh <> lt rh <> rt Nil{}:
275 match alignment:
276
277def Aligned.project_right(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, alignment: Aligned(left, right, output), coordinate: Coordinate(output)) -> Coordinate(right):
278 match left right output:
279 case Nil{} Nil{} Nil{}:
280 Unit{}
281 case +lh <> +lt +rh <> +rt +oh <> +ot:
282 AStep{axis, tail} = alignment
283 CStep{index, rest} = coordinate
284 CStep{Axis.project_right(lh, rh, oh, axis, index), Aligned.project_right(lt, rt, ot, tail, rest)}
285 case Nil{} Nil{} oh <> ot:
286 match alignment:
287 case Nil{} rh <> rt Nil{}:
288 match alignment:
289 case lh <> lt Nil{} Nil{}:
290 match alignment:
291 case Nil{} rh <> rt oh <> ot:
292 match alignment:
293 case lh <> lt Nil{} oh <> ot:
294 match alignment:
295 case lh <> lt rh <> rt Nil{}:
296 match alignment:
297
298def pad(count: Nat, dims: List<&2, Nat>) -> List<&2, Nat>:
299 match count:
300 case 0n:
301 dims
302 case 1n+rest:
303 1n <> pad(rest, dims)
304
305def unpad(+count: Nat, +dims: List<&2, Nat>, coordinate: Coordinate(pad(count, dims))) -> Coordinate(dims):
306 match count:
307 case 0n:
308 coordinate
309 case 1n+rest:
310 CStep{index, tail} = coordinate
311 unpad(rest, dims, tail)
312
313type PadMinimal<-left: Nat, -right: Nat> is Data:
314 PLeftZero{zero: {0n == left : Nat}}
315 PRightZero{zero: {0n == right : Nat}}
316
317type Broadcast<-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>> is Data:
318 BAlign{left_padding: Nat, right_padding: Nat, minimal: PadMinimal<left_padding, right_padding>, aligned: Aligned(pad(left_padding, left), pad(right_padding, right), output)}
319
320def project_left(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Broadcast<left, right, output>, coordinate: Coordinate(output)) -> Coordinate(left):
321 BAlign{left_padding, right_padding, minimal, aligned} = alignment
322 unpad(left_padding, left, Aligned.project_left(pad(left_padding, left), pad(right_padding, right), output, aligned, coordinate))
323
324def project_right(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Broadcast<left, right, output>, coordinate: Coordinate(output)) -> Coordinate(right):
325 BAlign{left_padding, right_padding, minimal, aligned} = alignment
326 unpad(right_padding, right, Aligned.project_right(pad(left_padding, left), pad(right_padding, right), output, aligned, coordinate))
327
328def offset(+dims: List<&2, Nat>, values: List<&2, Nat>) -> Nat:
329 match dims values:
330 case head <> tail value <> rest:
331 Nat.add(Nat.mul(value, numel(tail)), offset(tail, rest))
332 case dims values:
333 0n
334
335def digits(+dims: List<&2, Nat>, +index: Nat) -> List<&2, Nat>:
336 match dims:
337 case Nil{}:
338 Nil{}
339 case head <> tail:
340 Nat.div(index, numel(tail)) <> digits(tail, Nat.mod(index, numel(tail)))
341
342def fit(dims: List<&2, Nat>, values: List<&2, Nat>) -> Bool:
343 match dims values:
344 case Nil{} Nil{}:
345 True{}
346 case dh <> dt vh <> vt:
347 Bool.and(Nat.is_lt(vh, dh), fit(dt, vt))
348 case Nil{} vh <> vt:
349 False{}
350 case dh <> dt Nil{}:
351 False{}
352
353
354
355type Row is Data:
356 Row{stride: Nat, left: Nat, right: Nat}
357
358def index_rows(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>) -> List<&2, Row>:
359 match left right output:
360 case lh <> lt rh <> rt oh <> ot:
361 Row{numel(ot), broadcast_index(lh, numel(lt)), broadcast_index(rh, numel(rt))} <> index_rows(lt, rt, ot)
362 case left right output:
363 Nil{}
364
365def index_table(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>) -> List<&2, Row>:
366 index_rows(pad(Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, left)), left), pad(Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, right)), right), output)
367
368def rank_dimensions(dims: List<&2, U32>) -> Nat:
369 match dims:
370 case Nil{}:
371 0n
372 case head <> tail:
373 1n+rank_dimensions(tail)
374
375def naturals(dims: List<&2, U32>) -> List<&2, Nat>:
376 match dims:
377 case Nil{}:
378 Nil{}
379 case head <> tail:
380 U32.to_nat(head) <> naturals(tail)
381
382def words(values: List<&2, Nat>) -> List<&2, U32>:
383 match values:
384 case Nil{}:
385 Nil{}
386 case head <> tail:
387 U32.from_nat(head) <> words(tail)
388
389def logical_numel(dims: List<&2, U32>) -> Nat:
390 numel(naturals(dims))
391
392
393
394
395def numel_words(dims: List<&2, U32>) -> Nat:
396 match dims:
397 case Nil{}:
398 1n
399 case head <> tail:
400 Nat.mul(U32.to_nat(head), numel_words(tail))
401
402def horner(dims: List<&2, U32>, values: List<&2, U32>, acc: Nat) -> Nat:
403 match dims values:
404 case head <> tail value <> rest:
405 horner(tail, rest, Nat.add(Nat.mul(acc, U32.to_nat(head)), U32.to_nat(value)))
406 case dims values:
407 acc
408
409def digits_words(dims: List<&2, U32>, +index: Nat) -> List<&2, Nat>:
410 match dims:
411 case Nil{}:
412 Nil{}
413 case head <> +tail:
414 +stride = numel_words(tail)
415 Nat.div(index, stride) <> digits_words(tail, Nat.mod(index, stride))
416
417def fit_words(dims: List<&2, U32>, values: List<&2, U32>) -> Bool:
418 match dims values:
419 case Nil{} Nil{}:
420 True{}
421 case dh <> dt vh <> vt:
422 Bool.and(Nat.is_lt(U32.to_nat(vh), U32.to_nat(dh)), fit_words(dt, vt))
423 case Nil{} vh <> vt:
424 False{}
425 case dh <> dt Nil{}:
426 False{}
427
428def fit_values(dims: List<&2, U32>, values: List<&2, Nat>) -> Bool:
429 match dims values:
430 case Nil{} Nil{}:
431 True{}
432 case dh <> dt vh <> vt:
433 Bool.and(Nat.is_lt(vh, U32.to_nat(dh)), fit_values(dt, vt))
434 case Nil{} vh <> vt:
435 False{}
436 case dh <> dt Nil{}:
437 False{}
438
439def projected_words(dims: List<&2, U32>, values: List<&2, U32>) -> List<&2, Nat>:
440 match dims values:
441 case head <> tail value <> rest:
442 broadcast_index(U32.to_nat(head), U32.to_nat(value)) <> projected_words(tail, rest)
443 case dims values:
444 Nil{}
445
446def walk_left(rows: List<&2, Row>, +index: Nat) -> Nat:
447 match rows:
448 case Nil{}:
449 0n
450 case Row{+stride, left, right} <> rest:
451 Nat.add(Nat.mul(left, Nat.div(index, stride)), walk_left(rest, Nat.mod(index, stride)))
452
453def walk_right(rows: List<&2, Row>, +index: Nat) -> Nat:
454 match rows:
455 case Nil{}:
456 0n
457 case Row{+stride, left, right} <> rest:
458 Nat.add(Nat.mul(right, Nat.div(index, stride)), walk_right(rest, Nat.mod(index, stride)))
459
460def fit_values_spec(+dims: List<&2, U32>, +values: List<&2, Nat>) -> {fit_values(dims, values) == fit(naturals(dims), values) : Bool}:
461 match dims values:
462 case Nil{} Nil{}:
463 {==}
464 case +dh <> +dt +vh <> +vt:
465 Equal.cong(Bool, Bool, rest => Bool.and(Nat.is_lt(vh, U32.to_nat(dh)), rest), fit_values(dt, vt), fit(naturals(dt), vt), fit_values_spec(dt, vt))
466 case Nil{} vh <> vt:
467 {==}
468 case dh <> dt Nil{}:
469 {==}
470
471def fit_values_true(+dims: List<&2, U32>, +values: List<&2, Nat>, equal: {True{} == fit_values(dims, values) : Bool}) -> {True{} == fit(naturals(dims), values) : Bool}:
472 %fit_values_spec(dims, values) : {True{} == _ : Bool}
473 equal
474
475def index_rows_words(left: List<&2, U32>, right: List<&2, U32>, output: List<&2, U32>) -> List<&2, Row>:
476 match left right output:
477 case lh <> +lt rh <> +rt oh <> +ot:
478 Row{numel_words(ot), broadcast_index(U32.to_nat(lh), numel_words(lt)), broadcast_index(U32.to_nat(rh), numel_words(rt))} <> index_rows_words(lt, rt, ot)
479 case left right output:
480 Nil{}
481
482def numel_words_spec(+dims: List<&2, U32>) -> {numel_words(dims) == numel(naturals(dims)) : Nat}:
483 match dims:
484 case Nil{}:
485 {==}
486 case +head <> +tail:
487 Equal.cong(Nat, Nat, value => Nat.mul(U32.to_nat(head), value), numel_words(tail), numel(naturals(tail)), numel_words_spec(tail))
488
489def index_rows_spec(+left: List<&2, U32>, +right: List<&2, U32>, +output: List<&2, U32>) -> {index_rows_words(left, right, output) == index_rows(naturals(left), naturals(right), naturals(output)) : List<&2, Row>}:
490 match left right output:
491 case +lh <> +lt +rh <> +rt +oh <> +ot:
492 %Equal.sym(Nat, numel_words(ot), numel(naturals(ot)), numel_words_spec(ot)) : {Row{_, broadcast_index(U32.to_nat(lh), numel_words(lt)), broadcast_index(U32.to_nat(rh), numel_words(rt))} <> index_rows_words(lt, rt, ot) == index_rows(naturals(lh <> lt), naturals(rh <> rt), naturals(oh <> ot)) : List<&2, Row>}
493 %Equal.sym(Nat, numel_words(lt), numel(naturals(lt)), numel_words_spec(lt)) : {Row{numel(naturals(ot)), broadcast_index(U32.to_nat(lh), _), broadcast_index(U32.to_nat(rh), numel_words(rt))} <> index_rows_words(lt, rt, ot) == index_rows(naturals(lh <> lt), naturals(rh <> rt), naturals(oh <> ot)) : List<&2, Row>}
494 %Equal.sym(Nat, numel_words(rt), numel(naturals(rt)), numel_words_spec(rt)) : {Row{numel(naturals(ot)), broadcast_index(U32.to_nat(lh), numel(naturals(lt))), broadcast_index(U32.to_nat(rh), _)} <> index_rows_words(lt, rt, ot) == index_rows(naturals(lh <> lt), naturals(rh <> rt), naturals(oh <> ot)) : List<&2, Row>}
495 %Equal.sym(List<&2, Row>, index_rows_words(lt, rt, ot), index_rows(naturals(lt), naturals(rt), naturals(ot)), index_rows_spec(lt, rt, ot)) : {Row{numel(naturals(ot)), broadcast_index(U32.to_nat(lh), numel(naturals(lt))), broadcast_index(U32.to_nat(rh), numel(naturals(rt)))} <> _ == index_rows(naturals(lh <> lt), naturals(rh <> rt), naturals(oh <> ot)) : List<&2, Row>}
496 {==}
497 case Nil{} right output:
498 {==}
499 case lh <> lt Nil{} output:
500 {==}
501 case lh <> lt rh <> rt Nil{}:
502 {==}
503
504def rank_naturals(+dims: List<&2, U32>) -> {rank_dimensions(dims) == List.length(&2, Nat, naturals(dims)) : Nat}:
505 match dims:
506 case Nil{}:
507 {==}
508 case head <> tail:
509 Equal.cong(Nat, Nat, value => 1n+value, rank_dimensions(tail), List.length(&2, Nat, naturals(tail)), rank_naturals(tail))
510
511type BoolCase<-value: Bool> is Data:
512 BTrue{equal: {True{} == value : Bool}}
513 BFalse{equal: {False{} == value : Bool}}
514
515def bool_case(+value: Bool) -> BoolCase<value>:
516 match value:
517 case True{}:
518 BTrue{{==}}
519 case False{}:
520 BFalse{{==}}
521
522def truth(value: Bool) -> Data:
523 match value:
524 case True{}:
525 Unit
526 case False{}:
527 Empty
528
529def clash(equal: {True{} == False{} : Bool}) -> Empty:
530 %equal : truth(_)
531 Unit{}
532
533def LE.decided(+a: Nat, +b: Nat, decided: {True{} == Nat.is_le(a, b) : Bool}) -> {Unit == LE(a, b) : Data}:
534 match a b:
535 case 0n b:
536 {==}
537 case 1n+p 0n:
538 Empty.absurd({Unit == Empty : Data}, clash(decided))
539 case 1n+p 1n+q:
540 LE.decided(p, q, decided)
541
542
543
544def LE.from(+a: Nat, +b: Nat, decided: {True{} == Nat.is_le(a, b) : Bool}) -> LE(a, b):
545 %LE.decided(a, b, decided) : _
546 Unit{}
547
548def equal_decided(+a: Nat, +b: Nat, decided: {True{} == Nat.is_eq(a, b) : Bool}) -> {a == b : Nat}:
549 match a b:
550 case 0n 0n:
551 {==}
552 case 0n 1n+q:
553 Empty.absurd({0n == 1n+q : Nat}, clash(decided))
554 case 1n+p 0n:
555 Empty.absurd({1n+p == 0n : Nat}, clash(decided))
556 case 1n+p 1n+q:
557 %equal_decided(p, q, decided) : {1n+p == 1n+_ : Nat}
558 {==}
559
560def equal_from(+a: Nat, +b: Nat, decided: {True{} == Nat.is_eq(a, b) : Bool}) -> {a == b : Nat}:
561 %equal_decided(a, b, decided) : {a == _ : Nat}
562 {==}
563
564type Capacity<-dims: List<&2, Nat>> is Data:
565 Within{size: Nat, exact: {size == numel(dims) : Nat}, bound: {True{} == Nat.is_le(size, 1048576n) : Bool}}
566 Beyond{}
567
568def capacity.bound(-dims: List<&2, Nat>, +size: Nat, exact: {size == numel(dims) : Nat}, bound: BoolCase<Nat.is_le(size, 1048576n)>) -> Capacity<dims>:
569 match bound:
570 case BFalse{equal}:
571 Beyond{}
572 case BTrue{equal}:
573 Within{size, exact, equal}
574
575def capacity.scale(+head: Nat, -tail: List<&2, Nat>, +size: Nat, exact: {size == numel(tail) : Nat}, small: BoolCase<Bool.or(Nat.is_eq(size, 0n), Nat.is_le(head, 1048576n))>) -> Capacity<head <> tail>:
576 match small:
577 case BFalse{equal}:
578 Beyond{}
579 case BTrue{equal}:
580 +product = Nat.mul(head, size)
581 capacity.bound(head <> tail, product, Equal.cong(Nat, Nat, value => Nat.mul(head, value), size, numel(tail), exact), bool_case(Nat.is_le(product, 1048576n)))
582
583def capacity.zero(+head: Nat, -tail: List<&2, Nat>) -> Capacity<head <> tail>:
584 match head:
585 case 0n:
586 Within{0n, {==}, {==}}
587 case 1n+p:
588 Beyond{}
589
590
591
592def capacity.step(+head: Nat, -tail: List<&2, Nat>, rest: Capacity<tail>) -> Capacity<head <> tail>:
593 match rest:
594 case Beyond{}:
595 capacity.zero(head, tail)
596 case Within{+size, exact, bound}:
597 capacity.scale(head, tail, size, exact, bool_case(Bool.or(Nat.is_eq(size, 0n), Nat.is_le(head, 1048576n))))
598
599def capacity(+dims: List<&2, Nat>) -> Capacity<dims>:
600 match dims:
601 case Nil{}:
602 Within{1n, {==}, {==}}
603 case +head <> tail:
604 capacity.step(head, tail, capacity(tail))
605
606type Error is Data:
607 RankExceeded{}
608 CapacityExceeded{}
609 CoordinateRankMismatch{}
610 CoordinateOutOfBounds{}
611 FlatIndexOutOfBounds{}
612 IncompatibleBroadcast{}
613
614type Shape is Data:
615 SShape{dimensions: List<&2, U32>, size: Nat, exact: {size == logical_numel(dimensions) : Nat}, rank_bound: LE(rank_dimensions(dimensions), 4n), capacity_bound: LE(logical_numel(dimensions), 1048576n)}
616
617def dimensions(shape: Shape) -> List<&2, U32>:
618 SShape{dims, size, exact, rank_bound, capacity_bound} = shape
619 dims
620
621def elements(shape: Shape) -> Nat:
622 SShape{dims, size, exact, rank_bound, capacity_bound} = shape
623 size
624
625def extents(shape: Shape) -> List<&2, Nat>:
626 naturals(dimensions(shape))
627
628def elements_exact(+shape: Shape) -> {elements(shape) == numel(extents(shape)) : Nat}:
629 match shape:
630 case SShape{dims, size, exact, rank_bound, capacity_bound}:
631 exact
632
633def rank(shape: Shape) -> U32:
634 U32.from_nat(rank_dimensions(dimensions(shape)))
635
636def count(shape: Shape) -> U32:
637 U32.from_nat(elements(shape))
638
639type ConstructResult<-input: List<&2, U32>> is Data:
640 SError{error: Error}
641 SBuilt{shape: Shape, preserved: {dimensions(shape) == input : List<&2, U32>}}
642
643def construct.bound(-input: List<&2, U32>, +size: Nat, exact: {size == logical_numel(input) : Nat}, bound: {True{} == Nat.is_le(size, 1048576n) : Bool}) -> LE(logical_numel(input), 1048576n):
644 %exact : LE(_, 1048576n)
645 LE.from(size, 1048576n, bound)
646
647def construct.capacity(+input: List<&2, U32>, rank_bound: LE(rank_dimensions(input), 4n), capacity: Capacity<naturals(input)>) -> ConstructResult<input>:
648 match capacity:
649 case Beyond{}:
650 SError{CapacityExceeded{}}
651 case Within{+size, +exact, bound}:
652 SBuilt{SShape{input, size, exact, rank_bound, construct.bound(input, size, exact, bound)}, {==}}
653
654def construct.ranked(+input: List<&2, U32>, decision: BoolCase<Nat.is_le(rank_dimensions(input), 4n)>) -> ConstructResult<input>:
655 match decision:
656 case BFalse{equal}:
657 SError{RankExceeded{}}
658 case BTrue{equal}:
659 construct.capacity(input, LE.from(rank_dimensions(input), 4n, equal), capacity(naturals(input)))
660
661def construct(+input: List<&2, U32>) -> ConstructResult<input>:
662 construct.ranked(input, bool_case(Nat.is_le(rank_dimensions(input), 4n)))
663
664type CheckedCoordinate<-shape: Shape> is Data:
665 CValues{values: List<&2, Nat>, valid: {True{} == fit(extents(shape), values) : Bool}}
666
667def coordinate_values(-shape: Shape, coordinate: CheckedCoordinate<shape>) -> List<&2, U32>:
668 CValues{values, valid} = coordinate
669 words(values)
670
671def coordinate_checked.finish(+shape: Shape, +values: List<&2, Nat>, verdict: BoolCase<fit_values(dimensions(shape), values)>) -> Result<&2, &2, Error, CheckedCoordinate<shape>>:
672 match verdict:
673 case BFalse{equal}:
674 Fail{CoordinateOutOfBounds{}}
675 case BTrue{equal}:
676 Done{CValues{values, fit_values_true(dimensions(shape), values, equal)}}
677
678def coordinate_checked.nat(+shape: Shape, +values: List<&2, Nat>) -> Result<&2, &2, Error, CheckedCoordinate<shape>>:
679 coordinate_checked.finish(shape, values, bool_case(fit_values(dimensions(shape), values)))
680
681type FlatIndex<-shape: Shape> is Data:
682 FIndex{value: Nat, bound: {True{} == Nat.is_lt(value, elements(shape)) : Bool}}
683
684def index_value(-shape: Shape, index: FlatIndex<shape>) -> U32:
685 FIndex{value, bound} = index
686 U32.from_nat(value)
687
688def flat_checked.finish(-shape: Shape, +index: Nat, verdict: BoolCase<Nat.is_lt(index, elements(shape))>) -> Result<&2, &2, Error, FlatIndex<shape>>:
689 match verdict:
690 case BFalse{equal}:
691 Fail{FlatIndexOutOfBounds{}}
692 case BTrue{equal}:
693 Done{FIndex{index, equal}}
694
695def flat_checked(+shape: Shape, +index: Nat) -> Result<&2, &2, Error, FlatIndex<shape>>:
696 flat_checked.finish(shape, index, bool_case(Nat.is_lt(index, elements(shape))))
697
698def flatten_checked.fits(+shape: Shape, +values: List<&2, U32>, verdict: BoolCase<fit_words(dimensions(shape), values)>) -> Result<&2, &2, Error, FlatIndex<shape>>:
699 match verdict:
700 case BFalse{equal}:
701 Fail{CoordinateOutOfBounds{}}
702 case BTrue{equal}:
703 flat_checked(shape, horner(dimensions(shape), values, 0n))
704
705def flatten_checked.ranked(+shape: Shape, +values: List<&2, U32>, same_rank: BoolCase<Nat.is_eq(rank_dimensions(dimensions(shape)), rank_dimensions(values))>) -> Result<&2, &2, Error, FlatIndex<shape>>:
706 match same_rank:
707 case BFalse{equal}:
708 Fail{CoordinateRankMismatch{}}
709 case BTrue{equal}:
710 flatten_checked.fits(shape, values, bool_case(fit_words(dimensions(shape), values)))
711
712def flatten_checked(+shape: Shape, values: List<&2, U32>) -> Result<&2, &2, Error, FlatIndex<shape>>:
713 +copy = values
714 flatten_checked.ranked(shape, copy, bool_case(Nat.is_eq(rank_dimensions(dimensions(shape)), rank_dimensions(copy))))
715
716def unflatten_checked.finish(+shape: Shape, +index: Nat, verdict: BoolCase<Nat.is_lt(index, elements(shape))>) -> Result<&2, &2, Error, CheckedCoordinate<shape>>:
717 match verdict:
718 case BFalse{equal}:
719 Fail{FlatIndexOutOfBounds{}}
720 case BTrue{equal}:
721 coordinate_checked.nat(shape, digits_words(dimensions(shape), index))
722
723def unflatten_checked(+shape: Shape, index: U32) -> Result<&2, &2, Error, CheckedCoordinate<shape>>:
724 +value = U32.to_nat(index)
725 unflatten_checked.finish(shape, value, bool_case(Nat.is_lt(value, elements(shape))))
726
727type AxisChoice<-left: U32, -right: U32> is Data:
728 Picks{output: U32, axis: Axis<U32.to_nat(left), U32.to_nat(right), U32.to_nat(output)>}
729 Clashes{}
730
731def axis_choice.right_one(+left: U32, +right: U32, one: BoolCase<Nat.is_eq(1n, U32.to_nat(right))>) -> AxisChoice<left, right>:
732 match one:
733 case BFalse{equal}:
734 Clashes{}
735 case BTrue{equal}:
736 Picks{left, AxisRightOne{{==}, equal_from(1n, U32.to_nat(right), equal)}}
737
738def axis_choice.left_one(+left: U32, +right: U32, one: BoolCase<Nat.is_eq(1n, U32.to_nat(left))>) -> AxisChoice<left, right>:
739 match one:
740 case BFalse{equal}:
741 axis_choice.right_one(left, right, bool_case(Nat.is_eq(1n, U32.to_nat(right))))
742 case BTrue{equal}:
743 Picks{right, AxisLeftOne{equal_from(1n, U32.to_nat(left), equal), {==}}}
744
745def axis_choice.equal(+left: U32, +right: U32, same: BoolCase<Nat.is_eq(U32.to_nat(left), U32.to_nat(right))>) -> AxisChoice<left, right>:
746 match same:
747 case BFalse{equal}:
748 axis_choice.left_one(left, right, bool_case(Nat.is_eq(1n, U32.to_nat(left))))
749 case BTrue{equal}:
750 Picks{left, AxisEqual{{==}, equal_from(U32.to_nat(left), U32.to_nat(right), equal)}}
751
752def axis_choice(+left: U32, +right: U32) -> AxisChoice<left, right>:
753 axis_choice.equal(left, right, bool_case(Nat.is_eq(U32.to_nat(left), U32.to_nat(right))))
754
755type Alignment<-left: List<&2, U32>, -right: List<&2, U32>> is Data:
756 Aligns{output: List<&2, U32>, aligned: Aligned(naturals(left), naturals(right), naturals(output))}
757 Misaligned{}
758
759def align.finish(-lh: U32, -lt: List<&2, U32>, -rh: U32, -rt: List<&2, U32>, output: U32, axis: Axis<U32.to_nat(lh), U32.to_nat(rh), U32.to_nat(output)>, rest: Alignment<lt, rt>) -> Alignment<lh <> lt, rh <> rt>:
760 match rest:
761 case Misaligned{}:
762 Misaligned{}
763 case Aligns{tail, aligned}:
764 Aligns{output <> tail, AStep{axis, aligned}}
765
766def align.step(-lh: U32, -lt: List<&2, U32>, -rh: U32, -rt: List<&2, U32>, rest: Alignment<lt, rt>, choice: AxisChoice<lh, rh>) -> Alignment<lh <> lt, rh <> rt>:
767 match choice:
768 case Clashes{}:
769 Misaligned{}
770 case Picks{output, axis}:
771 align.finish(lh, lt, rh, rt, output, axis, rest)
772
773def align(+left: List<&2, U32>, +right: List<&2, U32>) -> Alignment<left, right>:
774 match left right:
775 case Nil{} Nil{}:
776 Aligns{Nil{}, Unit{}}
777 case +lh <> +lt +rh <> +rt:
778 align.step(lh, lt, rh, rt, align(lt, rt), axis_choice(lh, rh))
779 case Nil{} rh <> rt:
780 Misaligned{}
781 case lh <> lt Nil{}:
782 Misaligned{}
783
784def pad_words(count: Nat, dims: List<&2, U32>) -> List<&2, U32>:
785 match count:
786 case 0n:
787 dims
788 case 1n+rest:
789 1 <> pad_words(rest, dims)
790
791def naturals_pad(+count: Nat, -dims: List<&2, U32>) -> {naturals(pad_words(count, dims)) == pad(count, naturals(dims)) : List<&2, Nat>}:
792 match count:
793 case 0n:
794 {==}
795 case 1n+rest:
796 %naturals_pad(rest, dims) : {1n <> naturals(pad_words(rest, dims)) == 1n <> _ : List<&2, Nat>}
797 {==}
798
799def minimal(+left: Nat, +right: Nat) -> PadMinimal<Nat.sub(right, left), Nat.sub(left, right)>:
800 match left right:
801 case 0n 0n:
802 PLeftZero{{==}}
803 case 0n 1n+q:
804 PRightZero{{==}}
805 case 1n+p 0n:
806 PLeftZero{{==}}
807 case 1n+p 1n+q:
808 minimal(p, q)
809
810type Index<-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>> is Data:
811 Indexed{rows: List<&2, Row>, exact: {rows == index_table(left, right, output) : List<&2, Row>}}
812 Unindexed{zero: {0n == numel(output) : Nat}}
813
814def index_of.rows(+left: Shape, +right: Shape, +output: Shape, +rank: Nat) -> List<&2, Row>:
815 index_rows_words(pad_words(Nat.sub(rank, rank_dimensions(dimensions(left))), dimensions(left)), pad_words(Nat.sub(rank, rank_dimensions(dimensions(right))), dimensions(right)), dimensions(output))
816
817def index_of.exact(+left: Shape, +right: Shape, +output: Shape) -> {index_of.rows(left, right, output, rank_dimensions(dimensions(output))) == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}:
818 +lp = Nat.sub(rank_dimensions(dimensions(output)), rank_dimensions(dimensions(left)))
819 +rp = Nat.sub(rank_dimensions(dimensions(output)), rank_dimensions(dimensions(right)))
820 %Equal.sym(List<&2, Row>, index_rows_words(pad_words(lp, dimensions(left)), pad_words(rp, dimensions(right)), dimensions(output)), index_rows(naturals(pad_words(lp, dimensions(left))), naturals(pad_words(rp, dimensions(right))), extents(output)), index_rows_spec(pad_words(lp, dimensions(left)), pad_words(rp, dimensions(right)), dimensions(output))) : {_ == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}
821 %Equal.sym(List<&2, Nat>, naturals(pad_words(lp, dimensions(left))), pad(lp, extents(left)), naturals_pad(lp, dimensions(left))) : {index_rows(_, naturals(pad_words(rp, dimensions(right))), extents(output)) == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}
822 %Equal.sym(List<&2, Nat>, naturals(pad_words(rp, dimensions(right))), pad(rp, extents(right)), naturals_pad(rp, dimensions(right))) : {index_rows(pad(lp, extents(left)), _, extents(output)) == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}
823 %Equal.sym(Nat, rank_dimensions(dimensions(output)), List.length(&2, Nat, extents(output)), rank_naturals(dimensions(output))) : {index_rows(pad(Nat.sub(_, rank_dimensions(dimensions(left))), extents(left)), pad(Nat.sub(_, rank_dimensions(dimensions(right))), extents(right)), extents(output)) == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}
824 %Equal.sym(Nat, rank_dimensions(dimensions(left)), List.length(&2, Nat, extents(left)), rank_naturals(dimensions(left))) : {index_rows(pad(Nat.sub(List.length(&2, Nat, extents(output)), _), extents(left)), pad(Nat.sub(List.length(&2, Nat, extents(output)), rank_dimensions(dimensions(right))), extents(right)), extents(output)) == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}
825 %Equal.sym(Nat, rank_dimensions(dimensions(right)), List.length(&2, Nat, extents(right)), rank_naturals(dimensions(right))) : {index_rows(pad(Nat.sub(List.length(&2, Nat, extents(output)), List.length(&2, Nat, extents(left))), extents(left)), pad(Nat.sub(List.length(&2, Nat, extents(output)), _), extents(right)), extents(output)) == index_table(extents(left), extents(right), extents(output)) : List<&2, Row>}
826 {==}
827
828
829
830def index_of(+left: Shape, +right: Shape, +output: Shape, empty: BoolCase<Nat.is_eq(elements(output), 0n)>) -> Index<extents(left), extents(right), extents(output)>:
831 match empty:
832 case BTrue{equal}:
833 Unindexed{Equal.trans(Nat, 0n, elements(output), numel(extents(output)), Equal.sym(Nat, elements(output), 0n, equal_from(elements(output), 0n, equal)), elements_exact(output))}
834 case BFalse{equal}:
835 Indexed{index_of.rows(left, right, output, rank_dimensions(dimensions(output))), index_of.exact(left, right, output)}
836
837type BroadcastPlan is Data:
838 BPlan{left: Shape, right: Shape, output: Shape, alignment: Broadcast<extents(left), extents(right), extents(output)>}
839
840def left_shape(plan: BroadcastPlan) -> Shape:
841 BPlan{left, right, output, alignment} = plan
842 left
843
844def right_shape(plan: BroadcastPlan) -> Shape:
845 BPlan{left, right, output, alignment} = plan
846 right
847
848def output_shape(plan: BroadcastPlan) -> Shape:
849 BPlan{left, right, output, alignment} = plan
850 output
851
852def alignment(+plan: BroadcastPlan) -> Broadcast<extents(left_shape(plan)), extents(right_shape(plan)), extents(output_shape(plan))>:
853 match plan:
854 case BPlan{left, right, output, alignment}:
855 alignment
856
857def index_map(+plan: BroadcastPlan) -> Index<extents(left_shape(plan)), extents(right_shape(plan)), extents(output_shape(plan))>:
858 +output = output_shape(plan)
859 index_of(left_shape(plan), right_shape(plan), output, bool_case(Nat.is_eq(elements(output), 0n)))
860
861def broadcast_checked.realign(-left: List<&2, Nat>, -right: List<&2, Nat>, -shape: Shape, -output: List<&2, U32>, preserved: {dimensions(shape) == output : List<&2, U32>}, aligned: Aligned(left, right, naturals(output))) -> Aligned(left, right, extents(shape)):
862 %Equal.sym(List<&2, U32>, dimensions(shape), output, preserved) : Aligned(left, right, naturals(_))
863 aligned
864
865def broadcast_checked.construct(+left: Shape, +right: Shape, left_padding: Nat, right_padding: Nat, minimal: PadMinimal<left_padding, right_padding>, -output: List<&2, U32>, aligned: Aligned(pad(left_padding, extents(left)), pad(right_padding, extents(right)), naturals(output)), result: ConstructResult<output>) -> Result<&2, &2, Error, BroadcastPlan>:
866 match result:
867 case SError{error}:
868 Fail{error}
869 case SBuilt{shape, preserved}:
870 Done{BPlan{left, right, shape, BAlign{left_padding, right_padding, minimal, broadcast_checked.realign(pad(left_padding, extents(left)), pad(right_padding, extents(right)), shape, output, preserved, aligned)}}}
871
872def broadcast_checked.repad(+left_padding: Nat, +right_padding: Nat, -left: List<&2, U32>, -right: List<&2, U32>, -output: List<&2, U32>, aligned: Aligned(naturals(pad_words(left_padding, left)), naturals(pad_words(right_padding, right)), naturals(output))) -> Aligned(pad(left_padding, naturals(left)), pad(right_padding, naturals(right)), naturals(output)):
873 %naturals_pad(left_padding, left) : Aligned(_, pad(right_padding, naturals(right)), naturals(output))
874 %naturals_pad(right_padding, right) : Aligned(naturals(pad_words(left_padding, left)), _, naturals(output))
875 aligned
876
877def broadcast_checked.align(+left: Shape, +right: Shape, +left_padding: Nat, +right_padding: Nat, minimal: PadMinimal<left_padding, right_padding>, alignment: Alignment<pad_words(left_padding, dimensions(left)), pad_words(right_padding, dimensions(right))>) -> Result<&2, &2, Error, BroadcastPlan>:
878 match alignment:
879 case Misaligned{}:
880 Fail{IncompatibleBroadcast{}}
881 case Aligns{+output, aligned}:
882 broadcast_checked.construct(left, right, left_padding, right_padding, minimal, output, broadcast_checked.repad(left_padding, right_padding, dimensions(left), dimensions(right), output, aligned), construct(output))
883
884def broadcast_checked.ranked(+left: Shape, +right: Shape, +left_rank: Nat, +right_rank: Nat) -> Result<&2, &2, Error, BroadcastPlan>:
885 broadcast_checked.align(left, right, Nat.sub(right_rank, left_rank), Nat.sub(left_rank, right_rank), minimal(left_rank, right_rank), align(pad_words(Nat.sub(right_rank, left_rank), dimensions(left)), pad_words(Nat.sub(left_rank, right_rank), dimensions(right))))
886
887def broadcast_checked(+left: Shape, +right: Shape) -> Result<&2, &2, Error, BroadcastPlan>:
888 broadcast_checked.ranked(left, right, rank_dimensions(dimensions(left)), rank_dimensions(dimensions(right)))
889
890def projection_checked.fits(+input: Shape, -output: Shape, +values: List<&2, U32>, verdict: BoolCase<fit_words(dimensions(output), values)>) -> Result<&2, &2, Error, CheckedCoordinate<input>>:
891 match verdict:
892 case BFalse{equal}:
893 Fail{CoordinateOutOfBounds{}}
894 case BTrue{equal}:
895 coordinate_checked.nat(input, projected_words(dimensions(input), List.drop(&2, U32, values, Nat.sub(rank_dimensions(values), rank_dimensions(dimensions(input))))))
896
897def projection_checked.ranked(+input: Shape, +output: Shape, +values: List<&2, U32>, same_rank: BoolCase<Nat.is_eq(rank_dimensions(dimensions(output)), rank_dimensions(values))>) -> Result<&2, &2, Error, CheckedCoordinate<input>>:
898 match same_rank:
899 case BFalse{equal}:
900 Fail{CoordinateRankMismatch{}}
901 case BTrue{equal}:
902 projection_checked.fits(input, output, values, bool_case(fit_words(dimensions(output), values)))
903
904def projection_checked(+input: Shape, +output: Shape, values: List<&2, U32>) -> Result<&2, &2, Error, CheckedCoordinate<input>>:
905 +copy = values
906 projection_checked.ranked(input, output, copy, bool_case(Nat.is_eq(rank_dimensions(dimensions(output)), rank_dimensions(copy))))
907
908def project_left_checked(+plan: BroadcastPlan, values: List<&2, U32>) -> Result<&2, &2, Error, CheckedCoordinate<left_shape(plan)>>:
909 projection_checked(left_shape(plan), output_shape(plan), values)
910
911def project_right_checked(+plan: BroadcastPlan, values: List<&2, U32>) -> Result<&2, &2, Error, CheckedCoordinate<right_shape(plan)>>:
912 projection_checked(right_shape(plan), output_shape(plan), values)
913
914def error_code(error: Error) -> U32:
915 match error:
916 case RankExceeded{}:
917 1
918 case CapacityExceeded{}:
919 2
920 case CoordinateRankMismatch{}:
921 3
922 case CoordinateOutOfBounds{}:
923 4
924 case FlatIndexOutOfBounds{}:
925 5
926 case IncompatibleBroadcast{}:
927 6