src/tensor/observation.bend

bendtorch-core@0.1.0.0 · 447 lines · 75 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/layout.bend as Layout
5import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
6
7def dimensions_pair(pair: Core.Dynamic & Shape.Shape) -> Maybe<&2, List<&2, U32>>:
8  (dynamic, shape) = pair
9  Some{Shape.dimensions(shape)}
10
11def count_pair(pair: Core.Dynamic & U32) -> Maybe<&2, U32>:
12  (dynamic, count) = pair
13  Some{count}
14
15def dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
16  match result:
17    case Fail{error}:
18      None{}
19    case Done{tensor}:
20      dimensions_pair(Core.shape(Core.pack(shape, tensor)))
21
22def count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
23  match result:
24    case Fail{error}:
25      None{}
26    case Done{tensor}:
27      count_pair(Core.count(Core.pack(shape, tensor)))
28
29def point(-shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
30  match result:
31    case Fail{error}:
32      None{}
33    case Done{tensor}:
34      Core.read_value(shape, Core.read_flat(shape, tensor, index))
35
36def after_value_pair(-shape: Shape.Shape, +other: U32,
37  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Maybe<&2, F32>:
38  (tensor, result) = pair
39  Core.read_value(shape, Core.read_flat(shape, tensor, other))
40
41def after_value(-shape: Shape.Shape, +index: U32, +other: U32,
42  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
43  match result:
44    case Fail{error}:
45      None{}
46    case Done{tensor}:
47      after_value_pair(shape, other, Core.read_flat(shape, tensor, index))
48
49def after_length_pair(-shape: Shape.Shape,
50  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Nat:
51  (tensor, result) = pair
52  Core.logical_length(shape, tensor)
53
54def after_length(-shape: Shape.Shape, +index: U32,
55  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
56  match result:
57    case Fail{error}:
58      None{}
59    case Done{tensor}:
60      Some{after_length_pair(shape, Core.read_flat(shape, tensor, index))}
61
62def after_dimensions_pair(+shape: Shape.Shape,
63  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Maybe<&2, List<&2, U32>>:
64  (tensor, result) = pair
65  dimensions_pair(Core.shape(Core.pack(shape, tensor)))
66
67def after_dimensions(+shape: Shape.Shape, +index: U32,
68  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
69  match result:
70    case Fail{error}:
71      None{}
72    case Done{tensor}:
73      after_dimensions_pair(shape, Core.read_flat(shape, tensor, index))
74
75def after_count_pair(+shape: Shape.Shape,
76  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Maybe<&2, U32>:
77  (tensor, result) = pair
78  count_pair(Core.count(Core.pack(shape, tensor)))
79
80def after_count(+shape: Shape.Shape, +index: U32,
81  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
82  match result:
83    case Fail{error}:
84      None{}
85    case Done{tensor}:
86      after_count_pair(shape, Core.read_flat(shape, tensor, index))
87
88def write_done_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Bool:
89  (tensor, result) = pair
90  Result.is_done(&2, &2, Core.Error, Unit, result)
91
92def write_done(-shape: Shape.Shape, +index: U32, +value: F32,
93  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Bool>:
94  match result:
95    case Fail{error}:
96      None{}
97    case Done{tensor}:
98      Some{write_done_pair(shape, Core.write_flat(shape, tensor, index, value))}
99
100def written_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) ->
101  Maybe<&2, F32>:
102  (tensor, result) = pair
103  Core.read_value(shape, Core.read_flat(shape, tensor, other))
104
105def written_value(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
106  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
107  match result:
108    case Fail{error}:
109      None{}
110    case Done{tensor}:
111      written_value_pair(shape, other, Core.write_flat(shape, tensor, index, value))
112
113def written_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Nat:
114  (tensor, result) = pair
115  Core.logical_length(shape, tensor)
116
117def written_length(-shape: Shape.Shape, +index: U32, +value: F32,
118  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
119  match result:
120    case Fail{error}:
121      None{}
122    case Done{tensor}:
123      Some{written_length_pair(shape, Core.write_flat(shape, tensor, index, value))}
124
125def written_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) ->
126  Maybe<&2, List<&2, U32>>:
127  (tensor, result) = pair
128  dimensions_pair(Core.shape(Core.pack(shape, tensor)))
129
130def written_dimensions(+shape: Shape.Shape, +index: U32, +value: F32,
131  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
132  match result:
133    case Fail{error}:
134      None{}
135    case Done{tensor}:
136      written_dimensions_pair(shape, Core.write_flat(shape, tensor, index, value))
137
138def written_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Maybe<&2, U32>:
139  (tensor, result) = pair
140  count_pair(Core.count(Core.pack(shape, tensor)))
141
142def written_count(+shape: Shape.Shape, +index: U32, +value: F32,
143  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
144  match result:
145    case Fail{error}:
146      None{}
147    case Done{tensor}:
148      written_count_pair(shape, Core.write_flat(shape, tensor, index, value))
149
150def read_done_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Bool:
151  (tensor, result) = pair
152  Result.is_done(&2, &2, Core.Error, F32, result)
153
154def read_done(-shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Bool>:
155  match result:
156    case Fail{error}:
157      None{}
158    case Done{tensor}:
159      Some{read_done_pair(shape, Core.read_flat(shape, tensor, index))}
160
161def read_status(result: Result<&2, &2, Core.Error, F32>) -> U32:
162  match result:
163    case Fail{error}:
164      Core.error_code(error)
165    case Done{value}:
166      0
167
168def read_code_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> U32:
169  (tensor, result) = pair
170  read_status(result)
171
172def read_code(-shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
173  match result:
174    case Fail{error}:
175      None{}
176    case Done{tensor}:
177      Some{read_code_pair(shape, Core.read_flat(shape, tensor, index))}
178
179def write_status(result: Result<&2, &2, Core.Error, Unit>) -> U32:
180  match result:
181    case Fail{error}:
182      Core.error_code(error)
183    case Done{unit}:
184      0
185
186def write_code_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> U32:
187  (tensor, result) = pair
188  write_status(result)
189
190def write_code(-shape: Shape.Shape, +index: U32, +value: F32,
191  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
192  match result:
193    case Fail{error}:
194      None{}
195    case Done{tensor}:
196      Some{write_code_pair(shape, Core.write_flat(shape, tensor, index, value))}
197
198def coordinate_index(-shape: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>) -> Maybe<&2, U32>:
199  match result:
200    case Fail{error}:
201      None{}
202    case Done{index}:
203      Some{Shape.index_value(shape, index)}
204
205def coordinate_value_pair(+shape: Shape.Shape, +coordinates: List<&2, U32>,
206  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Maybe<&2, F32>:
207  (tensor, result) = pair
208  Core.read_value(shape, Core.read_coordinate(shape, tensor, coordinates))
209
210def coordinate_value(+shape: Shape.Shape, +index: U32, +value: F32, +coordinates: List<&2, U32>,
211  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
212  match result:
213    case Fail{error}:
214      None{}
215    case Done{tensor}:
216      coordinate_value_pair(shape, coordinates, Core.write_flat(shape, tensor, index, value))
217
218def coordinate_written(+shape: Shape.Shape, +coordinates: List<&2, U32>, +value: F32, +other: U32,
219  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
220  match result:
221    case Fail{error}:
222      None{}
223    case Done{tensor}:
224      written_value_pair(shape, other, Core.write_coordinate(shape, tensor, coordinates, value))
225
226def dynamic_point(+shape: Shape.Shape, +index: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
227  match result:
228    case Fail{error}:
229      None{}
230    case Done{tensor}:
231      Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, tensor}, index))
232
233def clone_first_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
234  (first, second) = pair
235  Core.read_value(shape, Core.read_flat(shape, first, other))
236
237def clone_second_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
238  (first, second) = pair
239  Core.read_value(shape, Core.read_flat(shape, second, other))
240
241def clone_first_value(-shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
242  match result:
243    case Fail{error}:
244      None{}
245    case Done{tensor}:
246      clone_first_value_pair(shape, other, Core.clone(shape, tensor))
247
248def clone_second_value(-shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
249  match result:
250    case Fail{error}:
251      None{}
252    case Done{tensor}:
253      clone_second_value_pair(shape, other, Core.clone(shape, tensor))
254
255def clone_first_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Nat:
256  (first, second) = pair
257  Core.logical_length(shape, first)
258
259def clone_second_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Nat:
260  (first, second) = pair
261  Core.logical_length(shape, second)
262
263def clone_first_length(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
264  match result:
265    case Fail{error}:
266      None{}
267    case Done{tensor}:
268      Some{clone_first_length_pair(shape, Core.clone(shape, tensor))}
269
270def clone_second_length(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
271  match result:
272    case Fail{error}:
273      None{}
274    case Done{tensor}:
275      Some{clone_second_length_pair(shape, Core.clone(shape, tensor))}
276
277def clone_first_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, List<&2, U32>>:
278  (first, second) = pair
279  dimensions_pair(Core.shape(Core.pack(shape, first)))
280
281def clone_second_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, List<&2, U32>>:
282  (first, second) = pair
283  dimensions_pair(Core.shape(Core.pack(shape, second)))
284
285def clone_first_dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
286  match result:
287    case Fail{error}:
288      None{}
289    case Done{tensor}:
290      clone_first_dimensions_pair(shape, Core.clone(shape, tensor))
291
292def clone_second_dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
293  match result:
294    case Fail{error}:
295      None{}
296    case Done{tensor}:
297      clone_second_dimensions_pair(shape, Core.clone(shape, tensor))
298
299def clone_first_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, U32>:
300  (first, second) = pair
301  count_pair(Core.count(Core.pack(shape, first)))
302
303def clone_second_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, U32>:
304  (first, second) = pair
305  count_pair(Core.count(Core.pack(shape, second)))
306
307def clone_first_count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
308  match result:
309    case Fail{error}:
310      None{}
311    case Done{tensor}:
312      clone_first_count_pair(shape, Core.clone(shape, tensor))
313
314def clone_second_count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
315  match result:
316    case Fail{error}:
317      None{}
318    case Done{tensor}:
319      clone_second_count_pair(shape, Core.clone(shape, tensor))
320
321def clone_written_first_pair(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
322  pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
323  (first, second) = pair
324  written_value_pair(shape, other, Core.write_flat(shape, first, index, value))
325
326def clone_written_first(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
327  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
328  match result:
329    case Fail{error}:
330      None{}
331    case Done{tensor}:
332      clone_written_first_pair(shape, index, value, other, Core.clone(shape, tensor))
333
334def untouched(-shape: Shape.Shape, +other: U32, written: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>,
335  second: Core.Tensor<shape>) -> Maybe<&2, F32>:
336  (tensor, result) = written
337  Core.read_value(shape, Core.read_flat(shape, second, other))
338
339def clone_written_second_pair(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
340  pair: Core.Tensor<shape> & Core.Tensor<shape>) -> Maybe<&2, F32>:
341  (first, second) = pair
342  untouched(shape, other, Core.write_flat(shape, first, index, value), second)
343
344def clone_written_second(-shape: Shape.Shape, +index: U32, +value: F32, +other: U32,
345  result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
346  match result:
347    case Fail{error}:
348      None{}
349    case Done{tensor}:
350      clone_written_second_pair(shape, index, value, other, Core.clone(shape, tensor))
351
352def length_number(-shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Nat:
353  (tensor, n) = pair
354  n
355
356def length_value(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
357  match result:
358    case Fail{error}:
359      None{}
360    case Done{tensor}:
361      Some{length_number(shape, Core.length(shape, tensor))}
362
363def length_owner_value_pair(-shape: Shape.Shape, +other: U32, pair: Core.Tensor<shape> & Nat) -> Maybe<&2, F32>:
364  (tensor, n) = pair
365  Core.read_value(shape, Core.read_flat(shape, tensor, other))
366
367def length_owner_value(-shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
368  match result:
369    case Fail{error}:
370      None{}
371    case Done{tensor}:
372      length_owner_value_pair(shape, other, Core.length(shape, tensor))
373
374def length_owner_length_pair(-shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Nat:
375  (tensor, n) = pair
376  Core.logical_length(shape, tensor)
377
378def length_owner_length(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, Nat>:
379  match result:
380    case Fail{error}:
381      None{}
382    case Done{tensor}:
383      Some{length_owner_length_pair(shape, Core.length(shape, tensor))}
384
385def length_owner_dimensions_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Maybe<&2, List<&2, U32>>:
386  (tensor, n) = pair
387  dimensions_pair(Core.shape(Core.pack(shape, tensor)))
388
389def length_owner_dimensions(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, List<&2, U32>>:
390  match result:
391    case Fail{error}:
392      None{}
393    case Done{tensor}:
394      length_owner_dimensions_pair(shape, Core.length(shape, tensor))
395
396def length_owner_count_pair(+shape: Shape.Shape, pair: Core.Tensor<shape> & Nat) -> Maybe<&2, U32>:
397  (tensor, n) = pair
398  count_pair(Core.count(Core.pack(shape, tensor)))
399
400def length_owner_count(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, U32>:
401  match result:
402    case Fail{error}:
403      None{}
404    case Done{tensor}:
405      length_owner_count_pair(shape, Core.length(shape, tensor))
406
407def dynamic_owner_value(+other: U32, dynamic: Core.Dynamic) -> Maybe<&2, F32>:
408  Core.dynamic_read_value(Core.read_dynamic(dynamic, other))
409
410def shape_owner_pair(+other: U32, pair: Core.Dynamic & Shape.Shape) -> Maybe<&2, F32>:
411  (dynamic, shape) = pair
412  dynamic_owner_value(other, dynamic)
413
414def shape_owner_value(+shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
415  match result:
416    case Fail{error}:
417      None{}
418    case Done{tensor}:
419      shape_owner_pair(other, Core.shape(Core.pack(shape, tensor)))
420
421def count_owner_pair(+other: U32, pair: Core.Dynamic & U32) -> Maybe<&2, F32>:
422  (dynamic, count) = pair
423  dynamic_owner_value(other, dynamic)
424
425def count_owner_value(+shape: Shape.Shape, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<shape>>) -> Maybe<&2, F32>:
426  match result:
427    case Fail{error}:
428      None{}
429    case Done{tensor}:
430      count_owner_pair(other, Core.count(Core.pack(shape, tensor)))
431
432def listed(values: List<&2, F32>, found: Maybe<&2, U32>) -> Maybe<&2, F32>:
433  match found:
434    case None{}:
435      None{}
436    case Some{index}:
437      List.get(&2, F32, values, U32.to_nat(index))
438
439def permute_source(+input: Shape.Shape, +axes: List<&2, U32>, +index: U32) -> Maybe<&2, U32>:
440  Layout.reference_source_coordinates(Shape.dimensions(input),
441    Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(input)), Shape.naturals(axes),
442      Shape.words(Shape.digits(Shape.naturals(Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))), U32.to_nat(index))), 0n))
443
444# The row-major source index that `permute_source` names for output index
445# `index`, as a natural (0 where it names none).
446def permute_source_index(+input: Shape.Shape, +axes: List<&2, U32>, +index: Nat) -> Nat:
447  U32.to_nat(Maybe.default(&2, U32, permute_source(input, axes, U32.from_nat(index)), 0))