src/tensor/transforms.bend
bendtorch-core@0.1.0.0 · 507 lines · 63 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/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
5import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
6import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
7import bendtorch-core@0.1.0.0/src/tensor/permute_build.bend as Build
8import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
9import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
10import bendtorch-core@0.1.0.0/src/storage/certificates.bend as Certificates
11
12type Error is Data:
13 LayoutError{error: Layout.Error}
14 TensorError{error: Core.Error}
15
16def elements_equal(+input: Shape.Shape, +target: Shape.Shape,
17 equal: {True{} == Nat.is_eq(Shape.elements(input), Shape.elements(target)) : Bool}) ->
18 {Shape.logical_numel(Shape.dimensions(input)) == Shape.logical_numel(Shape.dimensions(target)) : Nat}:
19 Equal.trans(Nat, Shape.logical_numel(Shape.dimensions(input)), Shape.elements(input), Shape.logical_numel(Shape.dimensions(target)),
20 Equal.sym(Nat, Shape.elements(input), Shape.logical_numel(Shape.dimensions(input)), Shape.elements_exact(input)),
21 Equal.trans(Nat, Shape.elements(input), Shape.elements(target), Shape.logical_numel(Shape.dimensions(target)),
22 Lifecycle.natural_equal(Shape.elements(input), Shape.elements(target), equal), Shape.elements_exact(target)))
23
24def reshape_equal(-input: Shape.Shape, +target: Shape.Shape, tensor: Core.Tensor<input>,
25 -same: {Shape.logical_numel(Shape.dimensions(input)) == Shape.logical_numel(Shape.dimensions(target)) : Nat}) -> Core.Dynamic:
26 Core.TTensor{storage, valid} = tensor
27 Core.pack(target, Core.TTensor{storage,
28 Equal.trans(Nat, Lifecycle.logical_length(storage), Shape.logical_numel(Shape.dimensions(input)),
29 Shape.logical_numel(Shape.dimensions(target)), valid, same)})
30
31def reshape_decided(+input: Shape.Shape, +target: Shape.Shape, tensor: Core.Tensor<input>,
32 decision: Shape.BoolCase<Nat.is_eq(Shape.elements(input), Shape.elements(target))>) ->
33 Result<&2, &1, Error, Core.Dynamic>:
34 match decision:
35 case Shape.BFalse{equal}:
36 Fail{LayoutError{Layout.ElementCountMismatch{}}}
37 case Shape.BTrue{equal}:
38 Done{reshape_equal(input, target, tensor,
39 elements_equal(input, target, equal))}
40
41def reshape_built(+input: Shape.Shape, tensor: Core.Tensor<input>, -dimensions: List<&2, U32>,
42 built: Shape.ConstructResult<dimensions>) -> Result<&2, &1, Error, Core.Dynamic>:
43 match built:
44 case Shape.SError{error}:
45 Fail{LayoutError{Layout.ShapeError{error}}}
46 case Shape.SBuilt{+target, preserved}:
47 reshape_decided(input, target, tensor,
48 Shape.bool_case(Nat.is_eq(Shape.elements(input), Shape.elements(target))))
49
50def reshape(input: Shape.Shape, tensor: Core.Tensor<input>, +dimensions: List<&2, U32>) -> Result<&2, &1, Error, Core.Dynamic>:
51 reshape_built(input, tensor, dimensions, Shape.construct(dimensions))
52
53def reshape_dynamic(dynamic: Core.Dynamic, dimensions: List<&2, U32>) -> Result<&2, &1, Error, Core.Dynamic>:
54 Core.Packed{input, tensor} = dynamic
55 reshape(input, tensor, dimensions)
56
57def copy_written(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
58 rec: Core.Tensor<input> -> Core.Tensor<Layout.output(input, axes, plan)> ->
59 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
60 source: Core.Tensor<input>, pair: Core.Tensor<Layout.output(input, axes, plan)> & Result<&2, &2, Core.Error, Unit>) ->
61 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
62 (target, result) = pair
63 match result:
64 case Fail{error}:
65 (source, Fail{TensorError{error}})
66 case Done{unit}:
67 rec(source)(target)
68
69def copy_read(-input: Shape.Shape, -axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, i: U32,
70 rec: Core.Tensor<input> -> Core.Tensor<Layout.output(input, axes, plan)> ->
71 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
72 target: Core.Tensor<Layout.output(input, axes, plan)>, pair: Core.Tensor<input> & Result<&2, &2, Core.Error, F32>) ->
73 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
74 (source, result) = pair
75 match result:
76 case Fail{error}:
77 (source, Fail{TensorError{error}})
78 case Done{value}:
79 copy_written(input, axes, plan, rec, source,
80 Core.write_flat(Layout.output(input, axes, plan), target, i, value))
81
82def copy_mapped(-input: Shape.Shape, -axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +i: U32,
83 rec: Core.Tensor<input> -> Core.Tensor<Layout.output(input, axes, plan)> ->
84 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
85 source: Core.Tensor<input>, target: Core.Tensor<Layout.output(input, axes, plan)>,
86 mapped: Result<&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, i>>) ->
87 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
88 match mapped:
89 case Fail{error}:
90 (source, Fail{LayoutError{error}})
91 case Done{mapping}:
92 copy_read(input, axes, plan, i, rec, target,
93 Core.read_flat(input, source, Layout.source_index(input, axes, plan, i, mapping)))
94
95def copy_step(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
96 source: Core.Tensor<input>, target: Core.Tensor<Layout.output(input, axes, plan)>, +i: U32) ->
97 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
98 copy_mapped(input, axes, plan, i, source => target => (source, Done{target}), source, target,
99 Layout.map_index(input, axes, plan, i))
100
101def copy_next(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
102 rec: Core.Tensor<input> -> Core.Tensor<Layout.output(input, axes, plan)> ->
103 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
104 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>) ->
105 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
106 (source, result) = pair
107 match result:
108 case Fail{error}:
109 (source, Fail{error})
110 case Done{target}:
111 rec(source)(target)
112
113type NaturalCopy<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -cursor: Nat> is Type:
114 NaturalCopied{pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
115 -index: U32, -word_exact: {U32.to_nat(index) == cursor : Nat}, -candidate: {index == U32.from_nat(cursor) : U32}}
116 NaturalError{source: Core.Tensor<input>}
117
118type CursorExact<-cursor: Nat, -index: U32> is Data:
119 CursorExact{-equal: {U32.to_nat(index) == cursor : Nat}}
120
121def natural_result(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
122 -cursor: Nat, -index: U32, -word_exact: {U32.to_nat(index) == cursor : Nat}, -candidate: {index == U32.from_nat(cursor) : U32},
123 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>) ->
124 NaturalCopy<input, axes, plan, cursor>:
125 NaturalCopied{pair, index, word_exact, candidate}
126
127def natural_decided(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +cursor: Nat,
128 +index: U32, source: Core.Tensor<input>, target: Core.Tensor<Layout.output(input, axes, plan)>,
129 -candidate: {index == U32.from_nat(cursor) : U32},
130 decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) -> NaturalCopy<input, axes, plan, cursor>:
131 match decision:
132 case Shape.BFalse{equal}:
133 NaturalError{source}
134 case Shape.BTrue{equal}:
135 natural_result(input, axes, plan, cursor, index, Lifecycle.natural_equal(U32.to_nat(index), cursor, equal), candidate,
136 copy_step(input, axes, plan, source, target, index))
137
138def natural_step(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +cursor: Nat,
139 source: Core.Tensor<input>, target: Core.Tensor<Layout.output(input, axes, plan)>) -> NaturalCopy<input, axes, plan, cursor>:
140 +index = U32.from_nat(cursor)
141 natural_decided(input, axes, plan, cursor, index, source, target, {==},
142 Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
143
144def CursorValid(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -cursor: Nat,
145 result: NaturalCopy<input, axes, plan, cursor>) -> Data:
146 match result:
147 case NaturalCopied{pair, index, exact, candidate}:
148 CursorExact<cursor, index>
149 case NaturalError{source}:
150 Unit
151
152def natural_next(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -cursor: Nat,
153 rec: Core.Tensor<input> -> Core.Tensor<Layout.output(input, axes, plan)> ->
154 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
155 result: NaturalCopy<input, axes, plan, cursor>) ->
156 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
157 match result:
158 case NaturalCopied{pair, index, exact, candidate}:
159 copy_next(input, axes, plan, rec, pair)
160 case NaturalError{source}:
161 (source, Fail{LayoutError{Layout.MappingError{}}})
162
163def copy_loop(fuel: Nat, +input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
164 source: Core.Tensor<input>, target: Core.Tensor<Layout.output(input, axes, plan)>, +cursor: Nat) ->
165 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
166 match fuel:
167 case 0n:
168 (source, Done{target})
169 case 1n+p:
170 natural_next(input, axes, plan, cursor,
171 source => target => copy_loop(p, input, axes, plan, source, target, 1n+cursor),
172 natural_step(input, axes, plan, cursor, source, target))
173
174def copy_pack(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
175 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>) ->
176 Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
177 (source, result) = pair
178 match result:
179 case Fail{error}:
180 (source, Fail{error})
181 case Done{target}:
182 (source, Done{Core.pack(Layout.output(input, axes, plan), target)})
183
184
185
186
187
188
189def loop_allocated(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
190 source: Core.Tensor<input>, allocated: Result<&2, &1, Core.Error, Core.Tensor<Layout.output(input, axes, plan)>>) ->
191 Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
192 match allocated:
193 case Fail{error}:
194 (source, Fail{TensorError{error}})
195 case Done{target}:
196 copy_pack(input, axes, plan,
197 copy_loop(Shape.elements(Layout.output(input, axes, plan)), input, axes, plan, source, target, 0n))
198
199
200
201def built_cloned(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +depth: Nat, +length: Nat,
202 -valid: {Storage.bounds(depth, length) == True{} : Bool},
203 -exact: {length == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
204 +d: Nat, +n: Nat, -array: Array<F32>, pair: Array<F32> & Array<F32>, -same: {pair == Array.clone(F32, array) : Array<F32> & Array<F32>},
205 -balanced: Model.Balanced(F32, d, array), -ok: {Storage.bounds(d, n) == True{} : Bool},
206 -fits: {Lifecycle.logical_length(Storage.Owned{d, n, array, balanced, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
207 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
208 (kept, copy) = pair
209 (Core.TTensor{Storage.Owned{d, n, kept, Certificates.clone_first(F32, d, array, (kept, copy), same, balanced), ok}, fits},
210 Done{Core.TTensor{Storage.Owned{depth, length, Build.build(depth, Build.Request{input, axes}, length, copy),
211 Build.balanced(depth, Build.Request{input, axes}, length, copy), valid}, exact}})
212
213def built_owned(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +depth: Nat, +length: Nat,
214 -valid: {Storage.bounds(depth, length) == True{} : Bool},
215 -exact: {length == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
216 source: Core.Tensor<input>) -> Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
217 Core.TTensor{storage, fits} = source
218 match storage:
219 case Storage.Empty{}:
220 (Core.TTensor{Storage.Empty{}, fits}, Fail{TensorError{Core.StorageError{Storage.OutOfBounds{}}}})
221 case Storage.Owned{+d, +n, array, balanced, ok}:
222 built_cloned(input, axes, plan, depth, length, valid, exact, d, n, array, Array.clone(F32, array), {==}, balanced, ok, fits)
223
224
225def built(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
226 source: Core.Tensor<input>, target: Core.Tensor<Layout.output(input, axes, plan)>) ->
227 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>:
228 Core.TTensor{storage, exact} = target
229 match storage:
230 case Storage.Empty{}:
231 (source, Done{Core.TTensor{Storage.Empty{}, exact}})
232 case Storage.Owned{+depth, +length, array, balanced, valid}:
233 built_owned(input, axes, plan, depth, length, valid, exact, source)
234
235def copy_allocated(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
236 source: Core.Tensor<input>, allocated: Result<&2, &1, Core.Error, Core.Tensor<Layout.output(input, axes, plan)>>) ->
237 Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
238 match allocated:
239 case Fail{error}:
240 (source, Fail{TensorError{error}})
241 case Done{target}:
242 copy_pack(input, axes, plan, built(input, axes, plan, source, target))
243
244def copy_planned(+input: Shape.Shape, +axes: List<&2, U32>, source: Core.Tensor<input>,
245 planned: Result<&2, &2, Layout.Error, Layout.Plan<input, axes>>) -> Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
246 match planned:
247 case Fail{error}:
248 (source, Fail{LayoutError{error}})
249 case Done{+plan}:
250 copy_allocated(input, axes, plan, source, Core.new(Layout.output(input, axes, plan), 0.0))
251
252def permute_copy(+input: Shape.Shape, source: Core.Tensor<input>, +axes: List<&2, U32>) ->
253 Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
254 copy_planned(input, axes, source, Layout.permutation(input, axes))
255
256def transpose_decided(+input: Shape.Shape, source: Core.Tensor<input>, +a: U32, +b: U32,
257 decision: Shape.BoolCase<Transpose.valid(input, a, b)>) ->
258 Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
259 match decision:
260 case Shape.BFalse{equal}:
261 (source, Fail{LayoutError{Layout.AxisOutOfBounds{}}})
262 case Shape.BTrue{equal}:
263 permute_copy(input, source, Layout.transpose_axes(Shape.dimensions(input), a, b, 0))
264
265def transpose_copy(+input: Shape.Shape, source: Core.Tensor<input>, +a: U32, +b: U32) ->
266 Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>:
267 transpose_decided(input, source, a, b,
268 Shape.bool_case(Transpose.valid(input, a, b)))
269
270def dynamic_result(+input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>) ->
271 Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
272 (source, result) = pair
273 (Core.pack(input, source), result)
274
275def permute_dynamic(dynamic: Core.Dynamic, axes: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
276 Core.Packed{+input, source} = dynamic
277 dynamic_result(input, permute_copy(input, source, axes))
278
279def transpose_dynamic(dynamic: Core.Dynamic, a: U32, b: U32) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
280 Core.Packed{+input, source} = dynamic
281 dynamic_result(input, transpose_copy(input, source, a, b))
282
283def dynamic_observe(dynamic: Core.Dynamic, index: U32) -> Maybe<&2, F32>:
284 Core.Packed{shape, tensor} = dynamic
285 Core.observe(shape, tensor, index)
286
287def dynamic_dimensions(dynamic: Core.Dynamic) -> List<&2, U32>:
288 Core.Packed{shape, tensor} = dynamic
289 Shape.dimensions(shape)
290
291def copy_observe(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
292 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
293 index: U32) -> Maybe<&2, F32>:
294 (source, result) = pair
295 match result:
296 case Fail{error}:
297 None{}
298 case Done{target}:
299 Core.observe(Layout.output(input, axes, plan), target, index)
300
301type CopyMeaning<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -index: U32,
302 -mapping: Layout.Mapping<input, axes, plan, index>, -source: Core.Tensor<input>,
303 -target: Core.Tensor<Layout.output(input, axes, plan)>> is Data:
304 CopyMeaning{-meaning: {copy_observe(input, axes, plan, copy_step(input, axes, plan, source, target, index), index) ==
305 Core.checked_observe(input, source, Layout.source_index(input, axes, plan, index, mapping)) : Maybe<&2, F32>}}
306
307def expected_mapped(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -index: U32,
308 source: Core.Tensor<input>, result: Result<&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, index>>) -> Maybe<&2, F32>:
309 match result:
310 case Fail{error}:
311 None{}
312 case Done{mapping}:
313 Core.checked_observe(input, source, Layout.source_index(input, axes, plan, index, mapping))
314
315def expected_output(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
316 source: Core.Tensor<input>, +index: U32) -> Maybe<&2, F32>:
317 expected_mapped(input, axes, plan, index, source, Layout.map_index(input, axes, plan, index))
318
319type Prefix<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
320 -original: Core.Tensor<input>, -source: Core.Tensor<input>, -target: Core.Tensor<Layout.output(input, axes, plan)>, -cursor: Nat> is Data:
321 Prefix{-source_kept: @+other: U32 ->
322 {Core.checked_observe(input, source, other) == Core.checked_observe(input, original, other) : Maybe<&2, F32>},
323 -words: @+other: Nat -> @-earlier: {Nat.is_lt(other, cursor) == True{} : Bool} ->
324 {U32.to_nat(U32.from_nat(other)) == other : Nat},
325 -checked: @+other: Nat -> @-earlier: {Nat.is_lt(other, cursor) == True{} : Bool} ->
326 Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), target}, U32.from_nat(other)),
327 -written: @+other: Nat -> @-earlier: {Nat.is_lt(other, cursor) == True{} : Bool} ->
328 {Core.checked_observe(Layout.output(input, axes, plan), target, U32.from_nat(other)) ==
329 expected_output(input, axes, plan, original, U32.from_nat(other)) : Maybe<&2, F32>}}
330
331type Advanced<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
332 -original: Core.Tensor<input>, -before_source: Core.Tensor<input>, -after_source: Core.Tensor<input>,
333 -before_target: Core.Tensor<Layout.output(input, axes, plan)>, -after_target: Core.Tensor<Layout.output(input, axes, plan)>, -cursor: Nat> is Data:
334 Advanced{-word_exact: {U32.to_nat(U32.from_nat(cursor)) == cursor : Nat},
335 -source_kept: @+other: U32 ->
336 {Core.checked_observe(input, after_source, other) == Core.checked_observe(input, before_source, other) : Maybe<&2, F32>},
337 -current: {Core.checked_observe(Layout.output(input, axes, plan), after_target, U32.from_nat(cursor)) ==
338 expected_output(input, axes, plan, original, U32.from_nat(cursor)) : Maybe<&2, F32>},
339 -current_checked: Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), after_target}, U32.from_nat(cursor)),
340 -past_checked: @+other: Nat -> @-earlier: {Nat.is_lt(other, cursor) == True{} : Bool} ->
341 Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), after_target}, U32.from_nat(other)),
342 -frame: @+other: Nat -> @-earlier: {Nat.is_lt(other, cursor) == True{} : Bool} ->
343 {Core.checked_observe(Layout.output(input, axes, plan), after_target, U32.from_nat(other)) ==
344 Core.checked_observe(Layout.output(input, axes, plan), before_target, U32.from_nat(other)) : Maybe<&2, F32>}}
345
346def SameSource(-input: Shape.Shape, -original: Core.Tensor<input>, -retained: Core.Tensor<input>) -> Type:
347 @+other: U32 -> {Core.checked_observe(input, retained, other) == Core.checked_observe(input, original, other) : Maybe<&2, F32>}
348
349def StepMeaning(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, -index: U32,
350 -original_source: Core.Tensor<input>, -original_target: Core.Tensor<Layout.output(input, axes, plan)>,
351 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>) -> Type:
352 (source, result) = pair
353 match result:
354 case Fail{error}:
355 Unit
356 case Done{target}:
357 SameSource(input, original_source, source) &
358 (Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), target}, index) &
359 ({Core.checked_observe(Layout.output(input, axes, plan), target, index) ==
360 expected_output(input, axes, plan, original_source, index) : Maybe<&2, F32>} &
361 (@+other: U32 -> @checked: Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), original_target}, other) ->
362 @different: {Nat.is_eq(U32.to_nat(index), U32.to_nat(other)) == False{} : Bool} ->
363 Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), target}, other) &
364 {Core.checked_observe(Layout.output(input, axes, plan), target, other) ==
365 Core.checked_observe(Layout.output(input, axes, plan), original_target, other) : Maybe<&2, F32>})))
366
367type StepCertificate<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -index: U32,
368 -original_source: Core.Tensor<input>, -original_target: Core.Tensor<Layout.output(input, axes, plan)>> is Type:
369 StepCertificate{pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
370 exact: {pair == copy_step(input, axes, plan, original_source, original_target, index) :
371 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>},
372 -meaning: StepMeaning(input, axes, plan, index, original_source, original_target, pair)}
373
374def NaturalMeaning(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, -cursor: Nat,
375 -original_source: Core.Tensor<input>, -original_target: Core.Tensor<Layout.output(input, axes, plan)>,
376 result: NaturalCopy<input, axes, plan, cursor>) -> Type:
377 match result:
378 case NaturalError{source}:
379 Unit
380 case NaturalCopied{pair, index, exact, candidate}:
381 StepMeaning(input, axes, plan, index, original_source, original_target, pair)
382
383type NaturalCertificate<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -cursor: Nat,
384 -original_source: Core.Tensor<input>, -original_target: Core.Tensor<Layout.output(input, axes, plan)>> is Type:
385 NaturalCertificate{result: NaturalCopy<input, axes, plan, cursor>,
386 exact: {result == natural_step(input, axes, plan, cursor, original_source, original_target) : NaturalCopy<input, axes, plan, cursor>},
387 -meaning: NaturalMeaning(input, axes, plan, cursor, original_source, original_target, result)}
388
389def LoopMeaning(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -original: Core.Tensor<input>, -limit: Nat,
390 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>) -> Data:
391 (source, result) = pair
392 match result:
393 case Fail{error}:
394 Unit
395 case Done{target}:
396 Prefix<input, axes, plan, original, source, target, limit>
397
398type LoopCertificate<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -original: Core.Tensor<input>,
399 -fuel: Nat, -cursor: Nat, -source: Core.Tensor<input>, -target: Core.Tensor<Layout.output(input, axes, plan)>> is Type:
400 LoopCertificate{pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>,
401 exact: {pair == copy_loop(fuel, input, axes, plan, source, target, cursor) :
402 Core.Tensor<input> & Result<&2, &1, Error, Core.Tensor<Layout.output(input, axes, plan)>>},
403 -meaning: LoopMeaning(input, axes, plan, original, Nat.add(fuel, cursor), pair)}
404
405type WholeCopyModel<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -source: Core.Tensor<input>,
406 -target: Core.Tensor<Layout.output(input, axes, plan)>> is Data:
407 WholeCopyModel{-meaning: LoopMeaning(input, axes, plan, source,
408 Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
409 copy_loop(Shape.elements(Layout.output(input, axes, plan)), input, axes, plan, source, target, 0n))}
410
411type PermuteOutput<-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>,
412 -retained: Core.Tensor<input>, -output: Core.Dynamic> is Data:
413 PermuteOutput{-plan: Layout.Plan<input, axes>,
414 -dimensions: {dynamic_dimensions(output) == Shape.dimensions(Layout.output(input, axes, plan)) : List<&2, U32>},
415 -kept: SameSource(input, original, retained),
416 -words: @+other: Nat -> @-bound: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool} ->
417 {U32.to_nat(U32.from_nat(other)) == other : Nat},
418 -checked: @+other: Nat -> @-bound: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool} ->
419 Core.ObservationChecked(output, U32.from_nat(other)),
420 -observed: @+other: Nat -> @-bound: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool} ->
421 {Core.checked_observe_dynamic(output, U32.from_nat(other)) == expected_output(input, axes, plan, original, U32.from_nat(other)) : Maybe<&2, F32>},
422 -linear: @+other: Nat -> @-bound: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool} ->
423 {dynamic_observe(output, U32.from_nat(other)) == expected_output(input, axes, plan, original, U32.from_nat(other)) : Maybe<&2, F32>}}
424
425def PublicMeaning(-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>,
426 pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>) -> Data:
427 (retained, result) = pair
428 match result:
429 case Fail{error}:
430 Unit
431 case Done{output}:
432 PermuteOutput<input, axes, original, retained, output>
433
434type PublicModel<-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>,
435 -pair: Core.Tensor<input> & Result<&2, &1, Error, Core.Dynamic>> is Data:
436 PublicModel{-meaning: PublicMeaning(input, axes, original, pair)}
437
438type ReshapeOutput<-input: Shape.Shape, -original: Core.Tensor<input>, -dimensions: List<&2, U32>, -output: Core.Dynamic> is Data:
439 ReshapeOutput{-shape: Shape.Shape,
440 -dimensions_exact: {dynamic_dimensions(output) == dimensions : List<&2, U32>},
441 -shape_exact: {Shape.dimensions(shape) == dimensions : List<&2, U32>},
442 -count: {Shape.logical_numel(Shape.dimensions(input)) == Shape.logical_numel(dimensions) : Nat},
443 -linear: @+index: U32 -> {dynamic_observe(output, index) == Core.observe(input, original, index) : Maybe<&2, F32>}}
444
445def ReshapeMeaning(-input: Shape.Shape, -original: Core.Tensor<input>, -dimensions: List<&2, U32>,
446 result: Result<&2, &1, Error, Core.Dynamic>) -> Data:
447 match result:
448 case Fail{error}:
449 Unit
450 case Done{output}:
451 ReshapeOutput<input, original, dimensions, output>
452
453type ReshapeModel<-input: Shape.Shape, -original: Core.Tensor<input>, -dimensions: List<&2, U32>> is Data:
454 ReshapeModel{-meaning: ReshapeMeaning(input, original, dimensions, reshape(input, original, dimensions))}
455
456type RestoredOutput<-input: Shape.Shape, -original: Core.Tensor<input>, -output: Core.Dynamic> is Data:
457 RestoredOutput{-dimensions: {dynamic_dimensions(output) == Shape.dimensions(input) : List<&2, U32>},
458 -linear: @+index: U32 -> {dynamic_observe(output, index) == Core.observe(input, original, index) : Maybe<&2, F32>}}
459
460def RestoreMeaning(-input: Shape.Shape, -original: Core.Tensor<input>, result: Result<&2, &1, Error, Core.Dynamic>) -> Data:
461 match result:
462 case Fail{error}:
463 Unit
464 case Done{output}:
465 RestoredOutput<input, original, output>
466
467def ReshapeInverseMeaning(+input: Shape.Shape, -original: Core.Tensor<input>,
468 result: Result<&2, &1, Error, Core.Dynamic>) -> Data:
469 match result:
470 case Fail{error}:
471 Unit
472 case Done{output}:
473 RestoreMeaning(input, original, reshape_dynamic(output, Shape.dimensions(input)))
474
475type ReshapeInverseModel<-input: Shape.Shape, -original: Core.Tensor<input>, -dimensions: List<&2, U32>> is Data:
476 ReshapeInverseModel{-meaning: ReshapeInverseMeaning(input, original, reshape(input, original, dimensions))}
477
478type ReshapeSuccessModel<-input: Shape.Shape, -target: Shape.Shape, -original: Core.Tensor<input>> is Data:
479 ReshapeSuccessModel{-success: {Result.is_done(&2, &1, Error, Core.Dynamic,
480 reshape(input, original, Shape.dimensions(target))) == True{} : Bool}}
481
482def reshape_inverse_success(+input: Shape.Shape, result: Result<&2, &1, Error, Core.Dynamic>) -> Bool:
483 match result:
484 case Fail{error}:
485 False{}
486 case Done{output}:
487 Result.is_done(&2, &1, Error, Core.Dynamic, reshape_dynamic(output, Shape.dimensions(input)))
488
489type ReshapeInverseSuccessModel<-input: Shape.Shape, -target: Shape.Shape, -original: Core.Tensor<input>> is Data:
490 ReshapeInverseSuccessModel{-success: {reshape_inverse_success(input, reshape(input, original, Shape.dimensions(target))) == True{} : Bool}}
491
492def error_code(error: Error) -> U32:
493 match error:
494 case LayoutError{Layout.ShapeError{error}}:
495 Shape.error_code(error)
496 case LayoutError{Layout.AxesLengthMismatch{}}:
497 12
498 case LayoutError{Layout.AxisOutOfBounds{}}:
499 13
500 case LayoutError{Layout.DuplicateAxis{}}:
501 14
502 case LayoutError{Layout.ElementCountMismatch{}}:
503 15
504 case LayoutError{Layout.MappingError{}}:
505 16
506 case TensorError{error}:
507 Core.error_code(error)