tensor_shape.bend
stelliferous@0.0.3.0 · 102 lines · 0 definitions
1# Storage-independent extents. View derives this descriptor from its axes.
2import Base
3import stelliferous@0.0.3.0/math_checked_word.bend as CheckedWord
4
5type Shape is Data:
6 Shape{extents: List<&2,U32>}
7
8def rank(shape: Shape) -> Nat:
9 Shape{extents} = shape
10 List.length(&2,U32,extents)
11
12def extent(shape: Shape,axis: Nat) -> Maybe<&2,U32>:
13 Shape{extents} = shape
14 List.get(&2,U32,extents,axis)
15
16def same_extents(left: List<&2,U32>,right: List<&2,U32>) -> Bool:
17 match left right:
18 case Nil{} Nil{}: True{}
19 case Con{x,xs} Con{y,ys}: U32.is_eq(x,y) && same_extents(xs,ys)
20 case left right: False{}
21
22def same(left: Shape,right: Shape) -> Bool:
23 Shape{xs} = left
24 Shape{ys} = right
25 same_extents(xs,ys)
26
27def has_zero(extents: List<&2,U32>) -> Bool:
28 match extents:
29 case Nil{}: False{}
30 case Con{x,xs}: U32.is_zero(x) || has_zero(xs)
31
32def checked_multiply(left: U32,right: U32) -> Maybe<&2,U32>:
33 CheckedWord.multiply(left,right)
34
35def product(extents: List<&2,U32>,accumulator: Maybe<&2,U32>) -> Maybe<&2,U32>:
36 match extents accumulator:
37 case Nil{} result: result
38 case Con{x,xs} Some{value}: product(xs,checked_multiply(value,x))
39 case Con{x,xs} None{}: None{}
40
41def count_nonempty(empty: Bool,extents: List<&2,U32>) -> Maybe<&2,U32>:
42 match empty:
43 case True{}: Some{0}
44 case False{}: product(extents,Some{1})
45
46# Scalar [] has one element. Any zero extent makes an empty shape, even when
47# an earlier partial product would overflow. Nonempty overflow is rejected.
48def numel(shape: Shape) -> Maybe<&2,U32>:
49 Shape{+extents} = shape
50 count_nonempty(has_zero(extents),extents)
51
52# NumPy broadcasting of one axis: equal extents, or one of them is 1.
53def broadcast_extent(same: Bool,left_one: Bool,right_one: Bool,left: U32,right: U32) -> Maybe<&2,U32>:
54 match same left_one right_one:
55 case True{} left_one right_one: Some{left}
56 case False{} True{} right_one: Some{right}
57 case False{} False{} True{}: Some{left}
58 case False{} False{} False{}: None{}
59
60def prepend_extent(extent: Maybe<&2,U32>,rest: Maybe<&2,List<&2,U32>>) -> Maybe<&2,List<&2,U32>>:
61 match extent rest:
62 case Some{extent} Some{rest}: Some{extent <> rest}
63 case extent rest: None{}
64
65def broadcast_same_rank(left: List<&2,U32>,right: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
66 match left right:
67 case Nil{} Nil{}: Some{[]}
68 case Con{+x,xs} Con{+y,ys}:
69 prepend_extent(broadcast_extent(U32.is_eq(x,y),U32.is_eq(x,1),U32.is_eq(y,1),x,y),broadcast_same_rank(xs,ys))
70 case left right: None{}
71
72# The extents with leading axes of extent 1 added up to rank.
73def with_rank(+extents: List<&2,U32>,rank: Nat) -> List<&2,U32>:
74 List.append(&2,U32,List.replicate(U32,Nat.sub(rank,List.length(&2,U32,extents)),1),extents)
75
76# The common extents of two operands under NumPy broadcasting: the extents are
77# aligned at the last axis, a missing axis counts as 1, and on each axis the
78# extents are equal or one of them is 1.
79def broadcast(+left: List<&2,U32>,+right: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
80 broadcast_same_rank(with_rank(left,List.length(&2,U32,right)),with_rank(right,List.length(&2,U32,left)))
81
82def sum_when(valid: Bool,value: U32) -> Maybe<&2,U32>:
83 match valid:
84 case True{}: Some{value}
85 case False{}: None{}
86
87def checked_add(+left: U32,+right: U32) -> Maybe<&2,U32>:
88 sum_when(U32.is_le(right,(4294967295 - left : U32)),(left + right : U32))
89
90def checked_add_to(left: Maybe<&2,U32>,right: U32) -> Maybe<&2,U32>:
91 match left:
92 case Some{left}: checked_add(left,right)
93 case None{}: None{}
94
95# Each extent with before elements added in front and after elements behind.
96# The three lists have one entry per axis and no sum may wrap.
97def padded(extents: List<&2,U32>,before: List<&2,U32>,after: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
98 match extents before after:
99 case Nil{} Nil{} Nil{}: Some{[]}
100 case Con{extent,extents} Con{front,before} Con{back,after}:
101 prepend_extent(checked_add_to(checked_add(front,extent),back),padded(extents,before,after))
102 case extents before after: None{}