src/shape/LAWS.bend

bendtorch-core@0.1.0.0 · 237 lines · 34 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/flat.bend as Flat
4
5def built(-input: List<&2, U32>, result: Shape.ConstructResult<input>) -> Maybe<&2, List<&2, U32>>:
6  match result:
7    case Shape.SError{error}:
8      None{}
9    case Shape.SBuilt{shape, preserved}:
10      Some{Shape.dimensions(shape)}
11
12def index_result(-shape: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>) -> Maybe<&2, U32>:
13  match result:
14    case Fail{error}:
15      None{}
16    case Done{index}:
17      Some{Shape.index_value(shape, index)}
18
19def coordinate_result(-shape: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.CheckedCoordinate<shape>>) -> Maybe<&2, List<&2, U32>>:
20  match result:
21    case Fail{error}:
22      None{}
23    case Done{coordinate}:
24      Some{Shape.coordinate_values(shape, coordinate)}
25
26def planned_dimensions(left: Shape.Shape, right: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> List<&2, U32> & List<&2, U32>:
27  match result:
28    case Fail{error}:
29      (Shape.dimensions(left), Shape.dimensions(right))
30    case Done{+plan}:
31      (Shape.dimensions(Shape.left_shape(plan)), Shape.dimensions(Shape.right_shape(plan)))
32
33def broadcast_built(result: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Bool:
34  match result:
35    case Fail{error}:
36      False{}
37    case Done{plan}:
38      True{}
39
40def fin_of.next(-bound: Nat, index: Maybe<&2, Shape.Fin(bound)>) -> Maybe<&2, Shape.Fin(1n+bound)>:
41  match index:
42    case None{}:
43      None{}
44    case Some{previous}:
45      Some{Shape.FNext{previous}}
46
47def fin_of(+bound: Nat, +value: Nat) -> Maybe<&2, Shape.Fin(bound)>:
48  match bound value:
49    case 0n value:
50      None{}
51    case 1n+p 0n:
52      Some{Shape.FZero{}}
53    case 1n+p 1n+q:
54      fin_of.next(p, fin_of(p, q))
55
56def decode.step(-dh: Nat, -dt: List<&2, Nat>, tail: Maybe<&2, Shape.Coordinate(dt)>, head: Maybe<&2, Shape.Fin(dh)>) -> Maybe<&2, Shape.Coordinate(dh <> dt)>:
57  match tail:
58    case None{}:
59      None{}
60    case Some{rest}:
61      match head:
62        case None{}:
63          None{}
64        case Some{index}:
65          Some{Shape.CStep{index, rest}}
66
67def decode(+dims: List<&2, Nat>, +values: List<&2, Nat>) -> Maybe<&2, Shape.Coordinate(dims)>:
68  match dims values:
69    case Nil{} Nil{}:
70      Some{Unit{}}
71    case +dh <> +dt +vh <> +vt:
72      decode.step(dh, dt, decode(dt, vt), fin_of(dh, vh))
73    case Nil{} vh <> vt:
74      None{}
75    case dh <> dt Nil{}:
76      None{}
77
78def flatten_spec(+dims: List<&2, Nat>, coordinate: Maybe<&2, Shape.Coordinate(dims)>) -> Maybe<&2, U32>:
79  match coordinate:
80    case None{}:
81      None{}
82    case Some{coordinate}:
83      Some{U32.from_nat(Shape.row_major(dims, coordinate))}
84
85def unflatten_spec(+dims: List<&2, Nat>, index: Maybe<&2, Shape.Fin(Shape.numel(dims))>) -> Maybe<&2, List<&2, U32>>:
86  match index:
87    case None{}:
88      None{}
89    case Some{index}:
90      Some{Shape.words(Shape.indices(dims, Shape.unflatten(dims, index)))}
91
92def project_left_spec(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, coordinate: Maybe<&2, Shape.Coordinate(output)>) -> Maybe<&2, List<&2, U32>>:
93  match coordinate:
94    case None{}:
95      None{}
96    case Some{coordinate}:
97      Some{Shape.words(Shape.indices(left, Shape.project_left(left, right, output, alignment, coordinate)))}
98
99def project_right_spec(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, coordinate: Maybe<&2, Shape.Coordinate(output)>) -> Maybe<&2, List<&2, U32>>:
100  match coordinate:
101    case None{}:
102      None{}
103    case Some{coordinate}:
104      Some{Shape.words(Shape.indices(right, Shape.project_right(left, right, output, alignment, coordinate)))}
105
106def position(-shape: Shape.Shape, index: Shape.FlatIndex<shape>) -> Nat:
107  Shape.FIndex{value, bound} = index
108  value
109
110def project_left_coordinate(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, index: Maybe<&2, Shape.Fin(Shape.numel(output))>) -> Maybe<&2, Shape.Coordinate(left)>:
111  match index:
112    case None{}:
113      None{}
114    case Some{index}:
115      Some{Shape.project_left(left, right, output, alignment, Shape.unflatten(output, index))}
116
117def project_right_coordinate(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, index: Maybe<&2, Shape.Fin(Shape.numel(output))>) -> Maybe<&2, Shape.Coordinate(right)>:
118  match index:
119    case None{}:
120      None{}
121    case Some{index}:
122      Some{Shape.project_right(left, right, output, alignment, Shape.unflatten(output, index))}
123
124law row_major_bounded:
125  for +dims: List<&2, Nat>
126  for +coordinate: Shape.Coordinate(dims)
127  Shape.LT(Shape.row_major(dims, coordinate), Shape.numel(dims))
128
129law flatten_unflatten:
130  for +dims: List<&2, Nat>
131  for +index: Shape.Fin(Shape.numel(dims))
132  {Shape.flatten(dims, Shape.unflatten(dims, index)) == index : Shape.Fin(Shape.numel(dims))}
133
134law unflatten_flatten:
135  for +dims: List<&2, Nat>
136  for +coordinate: Shape.Coordinate(dims)
137  {Shape.unflatten(dims, Shape.flatten(dims, coordinate)) == coordinate : Shape.Coordinate(dims)}
138
139law flatten_row_major:
140  for +dims: List<&2, Nat>
141  for +coordinate: Shape.Coordinate(dims)
142  {Shape.Fin.value(Shape.numel(dims), Shape.flatten(dims, coordinate)) == Shape.row_major(dims, coordinate) : Nat}
143
144law broadcast_left_projection:
145  for +left: List<&2, Nat>
146  for +right: List<&2, Nat>
147  for +output: List<&2, Nat>
148  for +alignment: Shape.Broadcast<left, right, output>
149  for +coordinate: Shape.Coordinate(output)
150  {Shape.indices(left, Shape.project_left(left, right, output, alignment, coordinate)) == Shape.project_indices(left, Shape.indices(output, coordinate)) : List<&2, Nat>}
151
152law broadcast_right_projection:
153  for +left: List<&2, Nat>
154  for +right: List<&2, Nat>
155  for +output: List<&2, Nat>
156  for +alignment: Shape.Broadcast<left, right, output>
157  for +coordinate: Shape.Coordinate(output)
158  {Shape.indices(right, Shape.project_right(left, right, output, alignment, coordinate)) == Shape.project_indices(right, Shape.indices(output, coordinate)) : List<&2, Nat>}
159
160law construction_complete:
161  for +input: List<&2, U32>
162  for rank_bound: Shape.LE(Shape.rank_dimensions(input), 4n)
163  for capacity_bound: Shape.LE(Shape.logical_numel(input), 1048576n)
164  {built(input, Shape.construct(input)) == Some{input} : Maybe<&2, List<&2, U32>>}
165
166law flatten_checked_refines:
167  for +shape: Shape.Shape
168  for +values: List<&2, U32>
169  {index_result(shape, Shape.flatten_checked(shape, values)) == flatten_spec(Shape.extents(shape), decode(Shape.extents(shape), Shape.naturals(values))) : Maybe<&2, U32>}
170
171law unflatten_checked_refines:
172  for +shape: Shape.Shape
173  for +index: U32
174  {coordinate_result(shape, Shape.unflatten_checked(shape, index)) == unflatten_spec(Shape.extents(shape), fin_of(Shape.numel(Shape.extents(shape)), U32.to_nat(index))) : Maybe<&2, List<&2, U32>>}
175
176law broadcast_output_unique:
177  for +left: List<&2, Nat>
178  for +right: List<&2, Nat>
179  for +output: List<&2, Nat>
180  for +other: List<&2, Nat>
181  for +first: Shape.Broadcast<left, right, output>
182  for +second: Shape.Broadcast<left, right, other>
183  {output == other : List<&2, Nat>}
184
185law broadcast_checked_plan:
186  for +left: Shape.Shape
187  for +right: Shape.Shape
188  {planned_dimensions(left, right, Shape.broadcast_checked(left, right)) == (Shape.dimensions(left), Shape.dimensions(right)) : List<&2, U32> & List<&2, U32>}
189
190law broadcast_checked_complete:
191  for +left: Shape.Shape
192  for +right: Shape.Shape
193  for +output: List<&2, U32>
194  for +alignment: Shape.Broadcast<Shape.extents(left), Shape.extents(right), Shape.naturals(output)>
195  for +rank_bound: Shape.LE(Shape.rank_dimensions(output), 4n)
196  for +capacity_bound: Shape.LE(Shape.logical_numel(output), 1048576n)
197  {broadcast_built(Shape.broadcast_checked(left, right)) == True{} : Bool}
198
199law project_left_checked_refines:
200  for +plan: Shape.BroadcastPlan
201  for +values: List<&2, U32>
202  {coordinate_result(Shape.left_shape(plan), Shape.project_left_checked(plan, values)) == project_left_spec(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.alignment(plan), decode(Shape.extents(Shape.output_shape(plan)), Shape.naturals(values))) : Maybe<&2, List<&2, U32>>}
203
204law project_right_checked_refines:
205  for +plan: Shape.BroadcastPlan
206  for +values: List<&2, U32>
207  {coordinate_result(Shape.right_shape(plan), Shape.project_right_checked(plan, values)) == project_right_spec(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.alignment(plan), decode(Shape.extents(Shape.output_shape(plan)), Shape.naturals(values))) : Maybe<&2, List<&2, U32>>}
208
209law project_left_index_refines:
210  for +plan: Shape.BroadcastPlan
211  for +map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>
212  for +index: Shape.FlatIndex<Shape.output_shape(plan)>
213  {Some{Shape.index_value(Shape.left_shape(plan), Flat.project_left_index(plan, map, index))} == flatten_spec(Shape.extents(Shape.left_shape(plan)), project_left_coordinate(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.alignment(plan), fin_of(Shape.numel(Shape.extents(Shape.output_shape(plan))), position(Shape.output_shape(plan), index)))) : Maybe<&2, U32>}
214
215law project_right_index_refines:
216  for +plan: Shape.BroadcastPlan
217  for +map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>
218  for +index: Shape.FlatIndex<Shape.output_shape(plan)>
219  {Some{Shape.index_value(Shape.right_shape(plan), Flat.project_right_index(plan, map, index))} == flatten_spec(Shape.extents(Shape.right_shape(plan)), project_right_coordinate(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.alignment(plan), fin_of(Shape.numel(Shape.extents(Shape.output_shape(plan))), position(Shape.output_shape(plan), index)))) : Maybe<&2, U32>}
220
221# A coordinate whose rank differs from the shape's fails with
222# CoordinateRankMismatch, checked before any bound: [9] on [2, 3] is a rank
223# error, not an out-of-bounds error.
224law flatten_rank_error:
225  for +shape: Shape.Shape
226  for +values: List<&2, U32>
227  for differ: {Shape.rank_dimensions(Shape.dimensions(shape)) != Shape.rank_dimensions(values) : Nat}
228  {Shape.flatten_checked(shape, values) == Fail{Shape.CoordinateRankMismatch{}} : Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>}
229
230# A coordinate of the shape's rank with a component outside its extent fails
231# with CoordinateOutOfBounds.
232law flatten_bounds_error:
233  for +shape: Shape.Shape
234  for +values: List<&2, U32>
235  for same: {Shape.rank_dimensions(Shape.dimensions(shape)) == Shape.rank_dimensions(values) : Nat}
236  for outside: {Shape.fit_words(Shape.dimensions(shape), values) == False{} : Bool}
237  {Shape.flatten_checked(shape, values) == Fail{Shape.CoordinateOutOfBounds{}} : Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>}