src/tensor/layout.bend
bendtorch-core@0.1.0.0 · 414 lines · 52 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
4
5type Error is Data:
6 ShapeError{error: Shape.Error}
7 AxesLengthMismatch{}
8 AxisOutOfBounds{}
9 DuplicateAxis{}
10 ElementCountMismatch{}
11 MappingError{}
12
13def contains(+axis: U32, axes: List<&2, U32>) -> Bool:
14 match axes:
15 case Nil{}:
16 False{}
17 case +head <> tail:
18 Bool.or(U32.is_eq(axis, head), contains(axis, tail))
19
20def unique(axes: List<&2, U32>) -> Bool:
21 match axes:
22 case Nil{}:
23 True{}
24 case +head <> +tail:
25 Bool.and(Bool.not(contains(head, tail)), unique(tail))
26
27def bounded(axes: List<&2, U32>, +rank: U32) -> Bool:
28 match axes:
29 case Nil{}:
30 True{}
31 case head <> tail:
32 Bool.and(U32.is_lt(head, rank), bounded(tail, rank))
33
34def bijective(+shape: Shape.Shape, +axes: List<&2, U32>) -> Bool:
35 Bool.and(Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(shape))),
36 Bool.and(bounded(axes, Shape.rank(shape)), unique(axes)))
37
38def lookup(values: List<&2, U32>, index: Nat) -> Maybe<&2, U32>:
39 match values index:
40 case Nil{} _:
41 None{}
42 case head <> tail 0n:
43 Some{head}
44 case head <> tail 1n+p:
45 lookup(tail, p)
46
47def reorder_tail(head: U32, result: Maybe<&2, List<&2, U32>>) -> Maybe<&2, List<&2, U32>>:
48 match result:
49 case None{}:
50 None{}
51 case Some{tail}:
52 Some{head <> tail}
53
54def reorder_head(rec: U32 -> Maybe<&2, List<&2, U32>>, head: Maybe<&2, U32>) -> Maybe<&2, List<&2, U32>>:
55 match head:
56 case None{}:
57 None{}
58 case Some{value}:
59 rec(value)
60
61def reorder(+values: List<&2, U32>, axes: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
62 match axes:
63 case Nil{}:
64 Some{Nil{}}
65 case axis <> tail:
66 reorder_head(value => reorder_tail(value, reorder(values, tail)), lookup(values, U32.to_nat(axis)))
67
68def coordinate_for(axes: List<&2, U32>, coordinates: List<&2, U32>, +axis: U32) -> Maybe<&2, U32>:
69 match axes coordinates:
70 case Nil{} _:
71 None{}
72 case _ Nil{}:
73 None{}
74 case +head <> tail value <> rest:
75 Bool.pick(Maybe<&2, U32>, U32.is_eq(head, axis), Some{value}, coordinate_for(tail, rest, axis))
76
77def inverse_coordinates(dimensions: List<&2, U32>, +axes: List<&2, U32>,
78 +coordinates: List<&2, U32>, +axis: U32) -> Maybe<&2, List<&2, U32>>:
79 match dimensions:
80 case Nil{}:
81 Some{Nil{}}
82 case dimension <> tail:
83 reorder_head(value => reorder_tail(value, inverse_coordinates(tail, axes, coordinates, U32.inc(axis))),
84 coordinate_for(axes, coordinates, axis))
85
86type Plan<-input: Shape.Shape, -axes: List<&2, U32>> is Data:
87 Permutation{output: Shape.Shape,
88 valid: {True{} == bijective(input, axes) : Bool},
89 ordered: {Some{Shape.dimensions(output)} == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>},
90 count: {True{} == Nat.is_eq(Shape.elements(input), Shape.elements(output)) : Bool}}
91
92def output(-input: Shape.Shape, -axes: List<&2, U32>, plan: Plan<input, axes>) -> Shape.Shape:
93 Permutation{shape, valid, ordered, count} = plan
94 shape
95
96def plan_counted(+input: Shape.Shape, +axes: List<&2, U32>, +out: Shape.Shape,
97 valid: {True{} == bijective(input, axes) : Bool},
98 ordered: {Some{Shape.dimensions(out)} == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>},
99 decision: Shape.BoolCase<Nat.is_eq(Shape.elements(input), Shape.elements(out))>) ->
100 Result<&2, &2, Error, Plan<input, axes>>:
101 match decision:
102 case Shape.BFalse{equal}:
103 Fail{ElementCountMismatch{}}
104 case Shape.BTrue{equal}:
105 Done{Permutation{out, valid, ordered, equal}}
106
107def plan_built(+input: Shape.Shape, +axes: List<&2, U32>, -dimensions: List<&2, U32>,
108 valid: {True{} == bijective(input, axes) : Bool},
109 selected: {Some{dimensions} == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>},
110 built: Shape.ConstructResult<dimensions>) -> Result<&2, &2, Error, Plan<input, axes>>:
111 match built:
112 case Shape.SError{error}:
113 Fail{ShapeError{error}}
114 case Shape.SBuilt{+out, preserved}:
115 plan_counted(input, axes, out, valid,
116 Equal.trans(Maybe<&2, List<&2, U32>>, Some{Shape.dimensions(out)}, Some{dimensions},
117 reorder(Shape.dimensions(input), axes),
118 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, xs => Some{xs}, Shape.dimensions(out), dimensions, preserved), selected),
119 Shape.bool_case(Nat.is_eq(Shape.elements(input), Shape.elements(out))))
120
121def plan_selected(+input: Shape.Shape, +axes: List<&2, U32>, valid: {True{} == bijective(input, axes) : Bool},
122 selected: Maybe<&2, List<&2, U32>>,
123 same: {selected == reorder(Shape.dimensions(input), axes) : Maybe<&2, List<&2, U32>>}) -> Result<&2, &2, Error, Plan<input, axes>>:
124 match selected:
125 case None{}:
126 Fail{MappingError{}}
127 case Some{+dimensions}:
128 plan_built(input, axes, dimensions, valid, same, Shape.construct(dimensions))
129
130def plan_decided(+input: Shape.Shape, +axes: List<&2, U32>, decision: Shape.BoolCase<bijective(input, axes)>) ->
131 Result<&2, &2, Error, Plan<input, axes>>:
132 match decision:
133 case Shape.BFalse{equal}:
134 Fail{MappingError{}}
135 case Shape.BTrue{equal}:
136 plan_selected(input, axes, equal, reorder(Shape.dimensions(input), axes), {==})
137
138def plan_unique(+input: Shape.Shape, +axes: List<&2, U32>, decision: Shape.BoolCase<unique(axes)>) ->
139 Result<&2, &2, Error, Plan<input, axes>>:
140 match decision:
141 case Shape.BFalse{equal}:
142 Fail{DuplicateAxis{}}
143 case Shape.BTrue{equal}:
144 plan_decided(input, axes, Shape.bool_case(bijective(input, axes)))
145
146def plan_bounded(+input: Shape.Shape, +axes: List<&2, U32>, decision: Shape.BoolCase<bounded(axes, Shape.rank(input))>) ->
147 Result<&2, &2, Error, Plan<input, axes>>:
148 match decision:
149 case Shape.BFalse{equal}:
150 Fail{AxisOutOfBounds{}}
151 case Shape.BTrue{equal}:
152 plan_unique(input, axes, Shape.bool_case(unique(axes)))
153
154def plan_ranked(+input: Shape.Shape, +axes: List<&2, U32>,
155 decision: Shape.BoolCase<Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input)))>) ->
156 Result<&2, &2, Error, Plan<input, axes>>:
157 match decision:
158 case Shape.BFalse{equal}:
159 Fail{AxesLengthMismatch{}}
160 case Shape.BTrue{equal}:
161 plan_bounded(input, axes, Shape.bool_case(bounded(axes, Shape.rank(input))))
162
163def permutation(+input: Shape.Shape, +axes: List<&2, U32>) -> Result<&2, &2, Error, Plan<input, axes>>:
164 plan_ranked(input, axes,
165 Shape.bool_case(Nat.is_eq(Shape.rank_dimensions(axes), Shape.rank_dimensions(Shape.dimensions(input)))))
166
167def transpose_axes(dimensions: List<&2, U32>, +a: U32, +b: U32, +i: U32) -> List<&2, U32>:
168 match dimensions:
169 case Nil{}:
170 Nil{}
171 case dimension <> tail:
172 Bool.pick(U32, U32.is_eq(i, a), b, Bool.pick(U32, U32.is_eq(i, b), a, i)) <>
173 transpose_axes(tail, a, b, U32.inc(i))
174
175def word_same(n: Nat, a: Word(n), b: Word(n)) -> Bool:
176 match n:
177 case 0n:
178 True{}
179 case 1n+p:
180 WCon{ab, at} = a
181 WCon{bb, bt} = b
182 Bool.and(Bool.not(Bool.xor(ab, bb)), word_same(p, at, bt))
183
184def word_value_same(a: U32, b: U32) -> Bool:
185 match a b:
186 case U32{x} U32{y}:
187 word_same(32n, x, y)
188
189def words_same(a: List<&2, U32>, b: List<&2, U32>) -> Bool:
190 match a b:
191 case Nil{} Nil{}:
192 True{}
193 case ah <> at bh <> bt:
194 Bool.and(word_value_same(ah, bh), words_same(at, bt))
195 case _ _:
196 False{}
197
198def maybe_same(a: Maybe<&2, List<&2, U32>>, b: Maybe<&2, List<&2, U32>>) -> Bool:
199 match a b:
200 case None{} None{}:
201 True{}
202 case Some{xs} Some{ys}:
203 words_same(xs, ys)
204 case _ _:
205 False{}
206
207type WordsValidStep<-dimension: U32, -dt: List<&2, U32>, -coordinate: U32, -ct: List<&2, U32>> is Data:
208 WStep{bound: {True{} == U32.is_lt(coordinate, dimension) : Bool}, tail: WordsValid(dt, ct)}
209
210def WordsValid(dims: List<&2, U32>, coordinates: List<&2, U32>) -> Data:
211 match dims coordinates:
212 case Nil{} Nil{}:
213 Unit
214 case dh <> dt ch <> ct:
215 WordsValidStep<dh, dt, ch, ct>
216 case _ _:
217 Empty
218
219def words_finish(-dimension: U32, -dt: List<&2, U32>, -coordinate: U32, -ct: List<&2, U32>,
220 bound: {True{} == U32.is_lt(coordinate, dimension) : Bool}, result: Result<&2, &2, Shape.Error, WordsValid(dt, ct)>) ->
221 Result<&2, &2, Shape.Error, WordsValid(dimension <> dt, coordinate <> ct)>:
222 match result:
223 case Fail{error}:
224 Fail{error}
225 case Done{tail}:
226 Done{WStep{bound, tail}}
227
228def words_decided(+dimension: U32, -dt: List<&2, U32>, +coordinate: U32, -ct: List<&2, U32>,
229 rec: Unit -> Result<&2, &2, Shape.Error, WordsValid(dt, ct)>, decision: Shape.BoolCase<U32.is_lt(coordinate, dimension)>) ->
230 Result<&2, &2, Shape.Error, WordsValid(dimension <> dt, coordinate <> ct)>:
231 match decision:
232 case Shape.BFalse{equal}:
233 Fail{Shape.CoordinateOutOfBounds{}}
234 case Shape.BTrue{equal}:
235 words_finish(dimension, dt, coordinate, ct, equal, rec(Unit{}))
236
237def words_check(+dims: List<&2, U32>, +coordinates: List<&2, U32>) -> Result<&2, &2, Shape.Error, WordsValid(dims, coordinates)>:
238 match dims coordinates:
239 case Nil{} Nil{}:
240 Done{Unit{}}
241 case +dh <> +dt +ch <> +ct:
242 words_decided(dh, dt, ch, ct, unit => words_check(dt, ct), Shape.bool_case(U32.is_lt(ch, dh)))
243 case _ _:
244 Fail{Shape.CoordinateRankMismatch{}}
245
246type CheckedCoordinate<-shape: Shape.Shape> is Data:
247 CWord{values: List<&2, U32>, valid: WordsValid(Shape.dimensions(shape), values)}
248
249def coordinate_values(-shape: Shape.Shape, coordinate: CheckedCoordinate<shape>) -> List<&2, U32>:
250 CWord{values, valid} = coordinate
251 values
252
253type FlatIndex<-shape: Shape.Shape> is Data:
254 FIndex{value: U32, valid: {True{} == U32.is_lt(value, Shape.count(shape)) : Bool}}
255
256def index_value(-shape: Shape.Shape, index: FlatIndex<shape>) -> U32:
257 FIndex{value, valid} = index
258 value
259
260def flatten_words(dims: List<&2, U32>, coordinates: List<&2, U32>, accumulator: U32) -> U32:
261 U32.from_nat(Shape.horner(dims, coordinates, U32.to_nat(accumulator)))
262
263def reverse_words_go(values: List<&2, U32>, accumulator: List<&2, U32>) -> List<&2, U32>:
264 match values:
265 case Nil{}:
266 accumulator
267 case head <> tail:
268 reverse_words_go(tail, head <> accumulator)
269
270def reverse_words(values: List<&2, U32>) -> List<&2, U32>:
271 reverse_words_go(values, Nil{})
272
273def unflatten_words(reversed_dimensions: List<&2, U32>, +index: U32) -> List<&2, U32>:
274 reverse_words(Shape.words(Shape.digits_words(reverse_words(reversed_dimensions), U32.to_nat(index))))
275
276def reference_source_coordinates(+dimensions: List<&2, U32>, coordinates: Maybe<&2, List<&2, U32>>) -> Maybe<&2, U32>:
277 match coordinates:
278 case None{}:
279 None{}
280 case Some{values}:
281 Some{U32.from_nat(Shape.offset(Shape.naturals(dimensions), Shape.naturals(values)))}
282
283def reference_source(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, index: U32) -> Maybe<&2, U32>:
284 reference_source_coordinates(Shape.dimensions(input),
285 Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(input)), Shape.naturals(axes),
286 Shape.words(Shape.digits(Shape.extents(output(input, axes, plan)), U32.to_nat(index))), 0n))
287
288type Mapping<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32> is Data:
289 Mapped{output_index: FlatIndex<output(input, axes, plan)>,
290 output_coordinate: CheckedCoordinate<output(input, axes, plan)>,
291 input_coordinate: CheckedCoordinate<input>, index: FlatIndex<input>,
292 requested: {index_value(output(input, axes, plan), output_index) == i : U32},
293 decoded: {coordinate_values(output(input, axes, plan), output_coordinate) ==
294 reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
295 inverse: {Some{coordinate_values(input, input_coordinate)} ==
296 inverse_coordinates(Shape.dimensions(input), axes,
297 coordinate_values(output(input, axes, plan), output_coordinate), 0) : Maybe<&2, List<&2, U32>>},
298 flattened: {index_value(input, index) ==
299 flatten_words(Shape.dimensions(input), coordinate_values(input, input_coordinate), 0) : U32},
300 roundtrip: {True{} == maybe_same(reorder(coordinate_values(input, input_coordinate), axes),
301 Some{coordinate_values(output(input, axes, plan), output_coordinate)}) : Bool}}
302
303def source_index(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32,
304 mapping: Mapping<input, axes, plan, i>) -> U32:
305 Mapped{oi, oc, ic, index, requested, decoded, inverse, flattened, roundtrip} = mapping
306 index_value(input, index)
307
308def output_words(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32,
309 mapping: Mapping<input, axes, plan, i>) -> List<&2, U32>:
310 Mapped{oi, oc, ic, index, requested, decoded, inverse, flattened, roundtrip} = mapping
311 coordinate_values(output(input, axes, plan), oc)
312
313def input_words(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Plan<input, axes>, -i: U32,
314 mapping: Mapping<input, axes, plan, i>) -> List<&2, U32>:
315 Mapped{oi, oc, ic, index, requested, decoded, inverse, flattened, roundtrip} = mapping
316 coordinate_values(input, ic)
317
318def map_flat(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
319 oi: FlatIndex<output(input, axes, plan)>, oc: CheckedCoordinate<output(input, axes, plan)>,
320 +ic: CheckedCoordinate<input>,
321 requested: {index_value(output(input, axes, plan), oi) == i : U32},
322 decoded: {coordinate_values(output(input, axes, plan), oc) ==
323 reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
324 inverse: {Some{coordinate_values(input, ic)} ==
325 inverse_coordinates(Shape.dimensions(input), axes, coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
326 roundtrip: {True{} == maybe_same(reorder(coordinate_values(input, ic), axes),
327 Some{coordinate_values(output(input, axes, plan), oc)}) : Bool},
328 +j: U32, flat: {j == flatten_words(Shape.dimensions(input), coordinate_values(input, ic), 0) : U32},
329 decision: Shape.BoolCase<U32.is_lt(j, Shape.count(input))>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
330 match decision:
331 case Shape.BFalse{equal}:
332 Fail{MappingError{}}
333 case Shape.BTrue{equal}:
334 Done{Mapped{oi, oc, ic, FIndex{j, equal}, requested, decoded, inverse, flat, roundtrip}}
335
336def map_roundtrip(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
337 oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
338 +ic: CheckedCoordinate<input>,
339 requested: {index_value(output(input, axes, plan), oi) == i : U32},
340 decoded: {coordinate_values(output(input, axes, plan), oc) ==
341 reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
342 inverse: {Some{coordinate_values(input, ic)} ==
343 inverse_coordinates(Shape.dimensions(input), axes, coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
344 decision: Shape.BoolCase<maybe_same(reorder(coordinate_values(input, ic), axes),
345 Some{coordinate_values(output(input, axes, plan), oc)})>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
346 match decision:
347 case Shape.BFalse{equal}:
348 Fail{MappingError{}}
349 case Shape.BTrue{equal}:
350 +j = flatten_words(Shape.dimensions(input), coordinate_values(input, ic), 0)
351 map_flat(input, axes, plan, i, oi, oc, ic, requested, decoded, inverse, equal, j, {==},
352 Shape.bool_case(U32.is_lt(j, Shape.count(input))))
353
354def map_input(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
355 oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
356 requested: {index_value(output(input, axes, plan), oi) == i : U32},
357 decoded: {coordinate_values(output(input, axes, plan), oc) ==
358 reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
359 +words: List<&2, U32>,
360 inverse: {Some{words} == inverse_coordinates(Shape.dimensions(input), axes,
361 coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
362 checked: Result<&2, &2, Shape.Error, WordsValid(Shape.dimensions(input), words)>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
363 match checked:
364 case Fail{error}:
365 Fail{ShapeError{error}}
366 case Done{valid}:
367 +ic = {CWord{words, valid} : CheckedCoordinate<input>}
368 map_roundtrip(input, axes, plan, i, oi, oc, ic, requested, decoded, inverse,
369 Shape.bool_case(maybe_same(reorder(words, axes), Some{coordinate_values(output(input, axes, plan), oc)})))
370
371def map_inverse(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
372 oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
373 requested: {index_value(output(input, axes, plan), oi) == i : U32},
374 decoded: {coordinate_values(output(input, axes, plan), oc) ==
375 reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
376 result: Maybe<&2, List<&2, U32>>,
377 same: {result == inverse_coordinates(Shape.dimensions(input), axes,
378 coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>}) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
379 match result:
380 case None{}:
381 Fail{MappingError{}}
382 case Some{+words}:
383 map_input(input, axes, plan, i, oi, oc, requested, decoded, words, same,
384 words_check(Shape.dimensions(input), words))
385
386def map_output(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
387 oi: FlatIndex<output(input, axes, plan)>,
388 requested: {index_value(output(input, axes, plan), oi) == i : U32},
389 +words: List<&2, U32>,
390 decoded: {words == reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
391 checked: Result<&2, &2, Shape.Error, WordsValid(Shape.dimensions(output(input, axes, plan)), words)>) ->
392 Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
393 match checked:
394 case Fail{error}:
395 Fail{ShapeError{error}}
396 case Done{valid}:
397 +oc = {CWord{words, valid} : CheckedCoordinate<output(input, axes, plan)>}
398 map_inverse(input, axes, plan, i, oi, oc, requested, decoded,
399 inverse_coordinates(Shape.dimensions(input), axes, words, 0), {==})
400
401def map_decided(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
402 decision: Shape.BoolCase<U32.is_lt(i, Shape.count(output(input, axes, plan)))>) ->
403 Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
404 match decision:
405 case Shape.BFalse{equal}:
406 Fail{ShapeError{Shape.FlatIndexOutOfBounds{}}}
407 case Shape.BTrue{equal}:
408 +words = reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i))
409 map_output(input, axes, plan, i, FIndex{i, equal}, {==}, words, {==},
410 words_check(Shape.dimensions(output(input, axes, plan)), words))
411
412def map_index(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32) ->
413 Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
414 map_decided(input, axes, plan, i, Shape.bool_case(U32.is_lt(i, Shape.count(output(input, axes, plan)))))