tensor_shape.bend

stelliferous@0.0.2.0 · 107 lines · 20 definitions

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