tensor_layout.bend
stelliferous@0.0.3.0 · 560 lines · 0 definitions
1# Layout operations consume their arguments. Copy keeps the bands and reshape
2# relabels them. Transpose, slice, flip, upsampling and set_slice go through a view
3# of the input: bands whose images and rows the view keeps in place run it on
4# their own rows, other inputs first arrive as one band.
5import Base
6import stelliferous@0.0.3.0/storage_buffer.bend as Storage
7import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
8import stelliferous@0.0.3.0/tensor_view.bend as Views
9import stelliferous@0.0.3.0/tensor_operations.bend as Operations
10import stelliferous@0.0.3.0/tensor.bend as Tensors
11import stelliferous@0.0.3.0/tensor_band.bend as Bands
12import stelliferous@0.0.3.0/traversal_partition.bend as Partition
13import stelliferous@0.0.3.0/tensor_band_concat.bend as BandConcat
14import stelliferous@0.0.3.0/tensor_band_schedule.bend as BandSchedule
15import stelliferous@0.0.3.0/tensor_band_view.bend as BandViews
16import stelliferous@0.0.3.0/tensor_band_reduce.bend as BandReduce
17
18def extents(view: Maybe<&2,Views.View>) -> List<&2,U32>:
19 match view:
20 case Some{view}: Views.axis_extents_of(view)
21 case None{}: []
22
23def naturals(values: List<&2,U32>) -> List<&2,Nat>:
24 match values:
25 case []: []
26 case value <> rest: U32.to_nat(value) <> naturals(rest)
27
28# The only layout operation that keeps its argument: (original, copy).
29def copy_bands(~Element: Data,status: Tensors.Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+fill: Element,+workers: U32) ->
30 Tensors.Tensor<Element> & Tensors.Tensor<Element>:
31 match status:
32 case Tensors.Usable{}:
33 Storage.bind_pair(~Bands.Bands(~Element),~Bands.Bands(~Element),~(Tensors.Tensor<Element> & Tensors.Tensor<Element>),
34 Bands.copy(~Element,bands),left => right =>
35 (Tensors.Tensor{left,geometry,shape,fill,Tensors.Usable{},workers},Tensors.Tensor{right,geometry,shape,fill,Tensors.Usable{},workers}))
36 case Tensors.Failed{+problem}:
37 (Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{problem},workers},
38 Tensors.whole(~Element,Tensors.buffer(~Element,fill,0),shape,fill,Tensors.Failed{problem},workers))
39
40def copy(~Element: Data,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
41 Tensors.Tensor{bands,geometry,shape,fill,status,workers} = tensor
42 copy_bands(~Element,status,bands,geometry,shape,fill,workers)
43
44def permuted(view: Maybe<&2,Views.View>,order: List<&2,Nat>) -> Maybe<&2,Views.View>:
45 match view:
46 case Some{view}: Views.permute(view,order)
47 case None{}: None{}
48
49# The result's axis i is the source axis order[i].
50def reordered(~Element: Data,name: String,tensor: Tensors.Tensor<Element>,order: List<&2,Nat>) -> Tensors.Tensor<Element>:
51 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
52 Tensors.gather_extents(~Element,name,status,permuted(Tensors.view_of(shape),order),[],bands,geometry,shape,fill,workers)
53
54def transpose(~Element: Data,tensor: Tensors.Tensor<Element>,order: List<&2,U32>) -> Tensors.Tensor<Element>:
55 reordered(~Element,"transpose",tensor,naturals(order))
56
57def sliced(view: Maybe<&2,Views.View>,axis: U32,start: U32,count: U32,step: U32) -> Maybe<&2,Views.View>:
58 match view:
59 case Some{view}: Views.slice(view,U32.to_nat(axis),start,count,step)
60 case None{}: None{}
61
62# Elements start, start + step, ... (count of them) along one axis.
63def slice(~Element: Data,tensor: Tensors.Tensor<Element>,axis: U32,start: U32,count: U32,step: U32) -> Tensors.Tensor<Element>:
64 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
65 Tensors.gather_extents(~Element,"slice",status,sliced(Tensors.view_of(shape),axis,start,count,step),[],bands,geometry,shape,fill,workers)
66
67def same_count(left: Maybe<&2,U32>,right: Maybe<&2,U32>) -> Bool:
68 match left right:
69 case Some{left} Some{right}: U32.is_eq(left,right)
70 case left right: False{}
71
72# The bands keep their geometry and order under the target shape when it
73# names as many elements.
74def relabel_ready(~Element: Data,accepted: Bool,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,fill: Element,
75 workers: U32) -> Tensors.Tensor<Element>:
76 match accepted:
77 case True{}: Tensors.Tensor{bands,geometry,target,fill,Tensors.Usable{},workers}
78 case False{}: Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{Tensors.Problem{"reshape",Tensors.ShapeMismatch{shape,target}}},workers}
79
80def relabel(~Element: Data,accepted: Bool,current: Tensors.Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,
81 fill: Element,workers: U32) -> Tensors.Tensor<Element>:
82 match current:
83 case Tensors.Usable{}: relabel_ready(~Element,accepted,bands,geometry,shape,target,fill,workers)
84 case Tensors.Failed{problem}: Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{problem},workers}
85
86# Same elements in the same row-major order under new extents.
87def reshape(~Element: Data,tensor: Tensors.Tensor<Element>,+shape: List<&2,U32>) -> Tensors.Tensor<Element>:
88 Tensors.Tensor{bands,+geometry,+source_shape,fill,status,workers} = tensor
89 relabel(~Element,same_count(Tensors.count_of(source_shape),Tensors.count_of(shape)),status,bands,geometry,source_shape,shape,fill,workers)
90
91def extent_at(extents: List<&2,U32>,axis: Nat) -> U32:
92 Tensors.known_count(List.get(&2,U32,extents,axis))
93
94def leading(ranked: Bool,+shape: List<&2,U32>) -> Maybe<&2,U32>:
95 match ranked:
96 case True{}: Shapes.product(List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),Some{1})
97 case False{}: None{}
98
99# The number of planes before the last two axes; none below rank two.
100def planes(+shape: List<&2,U32>) -> Maybe<&2,U32>:
101 leading(Nat.is_le(2n,List.length(&2,U32,shape)),shape)
102
103# Nearest-neighbour upsampling repeats each element of the last two axes
104# scale times, by copying from a view whose repeated axes have stride zero.
105def prepend_repeated(axis: Views.Axis,rest: Maybe<&2,List<&2,Views.Axis>>) -> Maybe<&2,List<&2,Views.Axis>>:
106 match rest:
107 case Some{axes}: Some{axis <> axes}
108 case None{}: None{}
109
110def repeated_axes(axes: List<&2,Views.Axis>,+scale: U32) -> Maybe<&2,List<&2,Views.Axis>>:
111 match axes:
112 case [Views.Axis{height,height_stride},Views.Axis{width,width_stride}]:
113 Some{[Views.Axis{height,height_stride},Views.Axis{scale,0},Views.Axis{width,width_stride},Views.Axis{scale,0}]}
114 case [axis] : None{}
115 case []: None{}
116 case axis <> rest: prepend_repeated(axis,repeated_axes(rest,scale))
117
118def with_offset(axes: Maybe<&2,List<&2,Views.Axis>>,offset: U32) -> Maybe<&2,Views.View>:
119 match axes:
120 case Some{axes}: Some{Views.View{axes,offset}}
121 case None{}: None{}
122
123def repeated(view: Maybe<&2,Views.View>,scale: U32) -> Maybe<&2,Views.View>:
124 match view:
125 case Some{Views.View{axes,offset}}: with_offset(repeated_axes(axes,scale),offset)
126 case None{}: None{}
127
128def scaled_extents(extents: List<&2,U32>,+scale: U32) -> List<&2,U32>:
129 match extents:
130 case [height,width]: [(height * scale : U32),(width * scale : U32)]
131 case extent <> rest: extent <> scaled_extents(rest,scale)
132 case []: []
133
134def upsample_checked(~Element: Data,valid: Bool,planes: Maybe<&2,U32>,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
135 +shape: List<&2,U32>,+scale: U32,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
136 match valid planes:
137 case True{} Some{planes}:
138 Tensors.gather(~Element,"upsample_nearest",Tensors.Usable{},repeated(Tensors.view_of(shape),scale),[],bands,geometry,shape,scaled_extents(shape,scale),fill,workers)
139 case valid planes: Tensors.failed(~Element,"upsample_nearest",Tensors.InvalidArgument{shape},scaled_extents(shape,scale),fill,workers)
140
141# The repeated extents and the output count fit the machine word.
142def upsample_valid(+shape: List<&2,U32>,+scale: U32) -> Bool:
143 +height = extent_at(shape,Nat.sub(List.length(&2,U32,shape),2n))
144 +width = extent_at(shape,Nat.sub(List.length(&2,U32,shape),1n))
145 Views.has_count(Shapes.checked_multiply(height,scale)) && Views.has_count(Shapes.checked_multiply(width,scale)) &&
146 Views.has_count(Tensors.count_of(scaled_extents(shape,scale)))
147
148def upsample_status(~Element: Data,status: Tensors.Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+scale: U32,
149 +fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
150 match status:
151 case Tensors.Failed{problem}: Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{problem},workers}
152 case Tensors.Usable{}: upsample_checked(~Element,upsample_valid(shape,scale),planes(shape),bands,geometry,shape,scale,fill,workers)
153
154# Nearest-neighbour upsampling of the last two axes: one gather through the
155# view that repeats each row and each column scale times.
156def upsample_nearest(~Element: Data,tensor: Tensors.Tensor<Element>,+scale: U32) -> Tensors.Tensor<Element>:
157 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
158 upsample_status(~Element,status,bands,geometry,shape,scale,fill,workers)
159
160def settle(~Element: Data,valid: Bool,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) ->
161 (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:
162 match result:
163 case Operations.Accepted{owners}: (owners,valid)
164 case Operations.Rejected{owners}: (owners,False{})
165
166# What every band of a placement shares: the banded slice of the target
167# geometry.
168type Placing is Data:
169 Placing{banded: BandViews.Banded,geometry: Bands.Geometry}
170
171def placed_outcome(~Element: Data,+first: U32,+rows: U32,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Bands.Transformed(~Element):
172 match result:
173 case Operations.Accepted{(part,target)}: Partition.Leaf{Bands.Checked{target,True{}},first,rows,Unit{}}
174 case Operations.Rejected{(part,target)}: Partition.Leaf{Bands.Checked{target,False{}},first,rows,Unit{}}
175
176def placed_viewed(~Element: Data,source: Maybe<&2,Views.View>,slice: Views.View,+first: U32,+rows: U32,target: Storage.Buffer<Element>,part: Storage.Buffer<Element>) ->
177 Bands.Transformed(~Element):
178 match source:
179 case Some{source}: placed_outcome(~Element,first,rows,Operations.scatter(~Element,part,source,target,slice))
180 case None{}: Partition.Leaf{Bands.Checked{target,False{}},first,rows,Unit{}}
181
182# A band of the target takes the band of the part of the same rows, written
183# through the local view of the slice.
184def placed_band(~Element: Data,context: Placing,first: U32,rows: U32,target: Storage.Buffer<Element>,part: Storage.Buffer<Element>) -> Bands.Transformed(~Element):
185 Placing{+banded,+geometry} = context
186 +rows = rows
187 placed_viewed(~Element,Tensors.view_of(BandViews.extents(banded,rows)),BandViews.local(banded,geometry,rows),first,rows,target,part)
188
189def place_planned(~Element: Data,+name: String,+banded: BandViews.Banded,+geometry: Bands.Geometry,+target_shape: List<&2,U32>,+fill: Element,+workers: U32,
190 part: Tensors.Tensor<Element>,owners: Bands.Bands(~Element) & Bands.Plan) -> Tensors.Tensor<Element>:
191 (target,+layout) = owners
192 Tensors.transformed(~Element,name,Tensors.InvalidArgument{target_shape},geometry,target_shape,fill,workers,
193 Bands.transform_paired(~Element,~Placing,~placed_band(~Element),layout,BandSchedule.parallel(geometry,workers,12.0,0.0,layout),Placing{banded,geometry},
194 Bands.zip(~Storage.Buffer<Element>,~Storage.Buffer<Element>,target,Tensors.arrive(~Element,BandViews.output(banded,geometry),layout,part))))
195
196def mismatched(~Element: Data,name: String,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32) ->
197 Tensors.Tensor<Element>:
198 Tensors.Tensor{bands,geometry,target_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.ShapeMismatch{target_shape,part_shape}}},workers}
199
200def place_whole(~Element: Data,+name: String,found: Maybe<&2,BandViews.Banded>,+whole: Bands.Geometry,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32,
201 part: Tensors.Tensor<Element>,target: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
202 match found:
203 case Some{+banded}: place_planned(~Element,name,banded,whole,target_shape,fill,workers,part,Bands.plan(~Element,Tensors.arrive(~Element,whole,Bands.unsplit(whole),target)))
204 case None{}: mismatched(~Element,name,Tensors.arrive(~Element,whole,Bands.unsplit(whole),target),whole,target_shape,part_shape,fill,workers)
205
206# The target first arrives as one band of one row.
207def place_target(~Element: Data,+name: String,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
208 +target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
209 +whole = {Bands.Geometry{1,1,1,Tensors.count_or_zero(target_shape)} : Bands.Geometry}
210 place_whole(~Element,name,BandViews.classify(axes,offset,whole),whole,target_shape,part_shape,fill,workers,part,Tensors.Tensor{bands,geometry,target_shape,fill,Tensors.Usable{},workers})
211
212# The part arrives in the rows of the target, so a slice writes band by band
213# only when it repeats no row.
214def place_rows(~Element: Data,+name: String,repeats: List<&2,U32>,+before: List<&2,Views.Axis>,+after: List<&2,Views.Axis>,+start: U32,+axes: List<&2,Views.Axis>,
215 +offset: U32,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32,
216 part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
217 match repeats:
218 case []: place_planned(~Element,name,BandViews.Banded{before,[],after,start},geometry,target_shape,fill,workers,part,Bands.plan(~Element,bands))
219 case extent <> rest: place_target(~Element,name,axes,offset,bands,geometry,target_shape,part_shape,fill,workers,part)
220
221# Target bands whose images and rows the slice keeps in place take the part
222# band by band; other targets first arrive as one band of one row.
223def place_found(~Element: Data,+name: String,found: Maybe<&2,BandViews.Banded>,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element),
224 +geometry: Bands.Geometry,
225 +target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
226 match found:
227 case Some{BandViews.Banded{+before,repeats,+after,+start}}:
228 place_rows(~Element,name,repeats,before,after,start,axes,offset,bands,geometry,target_shape,part_shape,fill,workers,part)
229 case None{}: place_target(~Element,name,axes,offset,bands,geometry,target_shape,part_shape,fill,workers,part)
230
231# A part of the extents of the slice is written; a part with no element
232# writes nothing.
233def place_checked(~Element: Data,+name: String,same: Bool,empty: Bool,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
234 +target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
235 match same empty:
236 case True{} True{}: Tensors.Tensor{bands,geometry,target_shape,fill,Tensors.Usable{},workers}
237 case True{} False{}: place_found(~Element,name,BandViews.classify(axes,offset,geometry),axes,offset,bands,geometry,target_shape,part_shape,fill,workers,part)
238 case False{} empty: mismatched(~Element,name,bands,geometry,target_shape,part_shape,fill,workers)
239
240# The slice comes taken apart, so that no View value is used twice: that
241# would make every View constructor shared.
242def place_slice(~Element: Data,+name: String,slice: Maybe<&2,Views.View>,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,
243 +fill: Element,+workers: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
244 match slice:
245 case Some{Views.View{+axes,+offset}}:
246 place_checked(~Element,name,Shapes.same_extents(Views.axis_extents(axes),part_shape),Shapes.has_zero(part_shape),axes,offset,bands,geometry,target_shape,part_shape,fill,
247 workers,part)
248 case None{}: mismatched(~Element,name,bands,geometry,target_shape,part_shape,fill,workers)
249
250def place_status(~Element: Data,+name: String,target_status: Tensors.Status,part_status: Tensors.Status,slice: Maybe<&2,Views.View>,bands: Bands.Bands(~Element),
251 +geometry: Bands.Geometry,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,+fill: Element,+workers: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
252 match target_status part_status:
253 case Tensors.Usable{} Tensors.Usable{}: place_slice(~Element,name,slice,bands,geometry,target_shape,part_shape,fill,workers,part)
254 case Tensors.Usable{} Tensors.Failed{+problem}: Tensors.Tensor{bands,geometry,target_shape,fill,Tensors.Failed{problem},workers}
255 case Tensors.Failed{problem} part_status: Tensors.Tensor{bands,geometry,target_shape,fill,Tensors.Failed{problem},workers}
256
257# target[..., start : start + part extent, ...] = part along axis: the part is
258# written through the slice of the target that it replaces.
259def set_slice(~Element: Data,target: Tensors.Tensor<Element>,+axis: U32,start: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
260 Tensors.Tensor{bands,+geometry,+target_shape,+fill,target_status,target_workers} = target
261 Tensors.Tensor{part_bands,part_geometry,+part_shape,part_fill,part_status,part_workers} = part
262 +workers = U32.max(target_workers,part_workers)
263 place_status(~Element,"set_slice",target_status,part_status,sliced(Tensors.view_of(target_shape),axis,start,extent_at(part_shape,U32.to_nat(axis)),1),bands,geometry,
264 target_shape,part_shape,fill,workers,Tensors.Tensor{part_bands,part_geometry,part_shape,part_fill,Tensors.Usable{},workers})
265
266# Concatenation along an axis.
267type Survey is Data:
268 Survey{valid: Bool,total: U32,status: Tensors.Status,workers: U32}
269
270def surveyed(survey: Survey,+first: List<&2,U32>,+shape: List<&2,U32>,+axis: Nat,status: Tensors.Status,workers: U32) -> Survey:
271 Survey{valid,+total,previous,most} = survey
272 +width = extent_at(shape,axis)
273 Survey{valid && Nat.is_eq(List.length(&2,U32,first),List.length(&2,U32,shape)) &&
274 Shapes.same_extents(List.take(&2,U32,first,axis),List.take(&2,U32,shape,axis)) &&
275 Shapes.same_extents(List.drop(&2,U32,first,1n+axis),List.drop(&2,U32,shape,1n+axis)) &&
276 U32.is_le(width,(4294967295 - total : U32)),
277 (total + width : U32),Tensors.first_failure(previous,status),U32.max(most,workers)}
278
279def prepend_survey(~Element: Data,tensor: Tensors.Tensor<Element>,result: Tensors.Parts<Element> & Survey) -> Tensors.Parts<Element> & Survey:
280 (rest,summary) = result
281 (Tensors.Cons{tensor,rest},summary)
282
283# Summarizes the extents of every part and returns the parts in order.
284def survey(~Element: Data,parts: Tensors.Parts<Element>,+first: List<&2,U32>,+axis: Nat,summary: Survey) -> Tensors.Parts<Element> & Survey:
285 match parts:
286 case Tensors.Single{tensor}:
287 Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & Survey),tensor,tensor => shape => fill => status => workers =>
288 +shape = shape
289 +status = status
290 +workers = workers
291 (Tensors.Single{tensor},surveyed(summary,first,shape,axis,status,workers)))
292 case Tensors.Cons{tensor,rest}:
293 Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & Survey),tensor,tensor => shape => fill => status => workers =>
294 +shape = shape
295 +status = status
296 +workers = workers
297 prepend_survey(~Element,tensor,survey(~Element,rest,first,axis,surveyed(summary,first,shape,axis,status,workers))))
298
299def with_extent(extents: List<&2,U32>,axis: Nat,+extent: U32) -> List<&2,U32>:
300 List.set(&2,U32,extents,axis,extent)
301
302def appended_shaped(~Element: Data,+axis: U32,+start: U32,output: Tensors.Tensor<Element>,shaped: Tensors.Tensor<Element> & List<&2,U32>) -> Tensors.Tensor<Element> & U32:
303 (part,+shape) = shaped
304 (set_slice(~Element,output,axis,start,part),(start + extent_at(shape,U32.to_nat(axis)) : U32))
305
306def append_at(~Element: Data,+axis: U32,+start: U32,output: Tensors.Tensor<Element>,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & U32:
307 appended_shaped(~Element,axis,start,output,Tensors.shape(~Element,part))
308
309def append_part(~Element: Data,+axis: U32,state: Tensors.Tensor<Element> & U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & U32:
310 (output,start) = state
311 append_at(~Element,axis,start,output,part)
312
313def appended_all(~Element: Data,state: Tensors.Tensor<Element> & U32) -> Tensors.Tensor<Element>:
314 (output,start) = state
315 output
316
317def append_parts(~Element: Data,parts: Tensors.Parts<Element>,+axis: U32,state: Tensors.Tensor<Element> & U32) -> Tensors.Tensor<Element>:
318 match parts:
319 case Tensors.Single{part}: appended_all(~Element,append_part(~Element,axis,state,part))
320 case Tensors.Cons{part,rest}: append_parts(~Element,rest,axis,append_part(~Element,axis,state,part))
321
322# The leaf concat of the bands of the parts at one band of the output, along
323# the middle axis of [outer, extent, inner], into a tensor of that band shape.
324def band_joined(~Element: Data,parts: Tensors.Parts<Element>,shape: List<&2,U32>,fill: Element) -> Tensors.Tensor<Element>:
325 append_parts(~Element,parts,1,(Tensors.with_workers(~Element,Tensors.full(~Element,shape,fill),1),0))
326
327# The parts meet in one plan of the rows of their joining, and every band of
328# the output appends the bands of the parts there.
329def join_counted(~Element: Data,count: Maybe<&2,U32>,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:
330 match count:
331 case Some{count}:
332 +joining = BandConcat.joining(shape,axis)
333 BandConcat.met(~Element,~band_joined(~Element),joining,shape,fill,workers,BandConcat.offers(~Element,joining,parts))
334 case None{}: Tensors.failed(~Element,"concat",Tensors.InvalidShape{shape},shape,fill,workers)
335
336# An empty result has no element to place.
337def join_shaped(~Element: Data,empty: Bool,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:
338 match empty:
339 case False{}: join_counted(~Element,Tensors.count_of(shape),shape,axis,fill,workers,parts)
340 case True{}: Tensors.with_workers(~Element,Tensors.full(~Element,shape,fill),workers)
341
342def join_status(~Element: Data,status: Tensors.Status,valid: Bool,+first: List<&2,U32>,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,
343 parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:
344 match status valid:
345 case Tensors.Usable{} True{}: join_shaped(~Element,Shapes.has_zero(shape),shape,axis,fill,workers,parts)
346 case Tensors.Usable{} False{}: Tensors.failed(~Element,"concat",Tensors.InvalidArgument{first},first,fill,workers)
347 case Tensors.Failed{problem} other: Tensors.whole(~Element,Tensors.buffer(~Element,fill,0),first,fill,Tensors.Failed{problem},workers)
348
349def joined(~Element: Data,+first: List<&2,U32>,+axis: U32,+fill: Element,surveyed: Tensors.Parts<Element> & Survey) -> Tensors.Tensor<Element>:
350 (parts,summary) = surveyed
351 Survey{valid,total,status,workers} = summary
352 join_status(~Element,status,valid,first,with_extent(first,U32.to_nat(axis),total),axis,fill,workers,parts)
353
354def lead(~Element: Data,parts: Tensors.Parts<Element>) -> Tensors.Parts<Element> & (List<&2,U32> & Element):
355 match parts:
356 case Tensors.Single{tensor}:
357 Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers =>
358 +shape = shape
359 +fill = fill
360 (Tensors.Single{tensor},(shape,fill)))
361 case Tensors.Cons{tensor,rest}:
362 Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers =>
363 +shape = shape
364 +fill = fill
365 (Tensors.Cons{tensor,rest},(shape,fill)))
366
367def concat_first(~Element: Data,+axis: U32,parts: Tensors.Parts<Element>,+shape: List<&2,U32>,+fill: Element) -> Tensors.Tensor<Element>:
368 joined(~Element,shape,axis,fill,
369 survey(~Element,parts,shape,U32.to_nat(axis),Survey{Nat.is_lt(U32.to_nat(axis),List.length(&2,U32,shape)),0,Tensors.Usable{},1}))
370
371def concat_led(~Element: Data,+axis: U32,led: Tensors.Parts<Element> & (List<&2,U32> & Element)) -> Tensors.Tensor<Element>:
372 (parts,first) = led
373 (shape,fill) = first
374 concat_first(~Element,axis,parts,shape,fill)
375
376# Parts agree on every extent except axis; the result stacks them in order.
377def concat(~Element: Data,parts: Tensors.Parts<Element>,+axis: U32) -> Tensors.Tensor<Element>:
378 concat_led(~Element,axis,lead(~Element,parts))
379
380# Every listed axis is below the rank.
381def all_below(axes: List<&2,U32>,+rank: Nat) -> Bool:
382 match axes:
383 case []: True{}
384 case axis <> rest: Nat.is_lt(U32.to_nat(axis),rank) && all_below(rest,rank)
385
386# Reduction over listed axes. Output k, at the row-major index of the kept
387# axes, folds the values of the reduced axes in row-major order, from
388# initial, and holds finish of the fold:
389# finish(operation(... operation(operation(initial, x[k, 0]), x[k, 1]) ..., x[k, R - 1])).
390# The axes after the kept ones, in increasing index, are the reduced ones.
391def chosen(member: Bool,+listed: Bool,+index: Nat,rest: List<&2,Nat>) -> List<&2,Nat>:
392 match member listed:
393 case True{} True{}: index <> rest
394 case False{} False{}: index <> rest
395 case member listed: rest
396
397def axes_where(remaining: Nat,+index: Nat,+axes: List<&2,Nat>,+listed: Bool) -> List<&2,Nat>:
398 match remaining:
399 case 0n: []
400 case 1n+rest: chosen(Views.contains_index(axes,index),listed,index,axes_where(rest,1n+index,axes,listed))
401
402def extents_at(+shape: List<&2,U32>,indices: List<&2,Nat>) -> List<&2,U32>:
403 match indices:
404 case []: []
405 case index <> rest: extent_at(shape,index) <> extents_at(shape,rest)
406
407def reduce_counted(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,+name: String,
408 +operations: F32,count: Maybe<&2,U32>,length: Maybe<&2,U32>,+initial: Accumulator,+shape: List<&2,U32>,+kept: List<&2,Nat>,+reduced: List<&2,Nat>,+fill: Element,
409 +workers: U32,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
410 match count length:
411 case Some{+count} Some{+length}:
412 BandReduce.lines(~Element,~Accumulator,~operation,~finish,name,operations,count,length,initial,extents_at(shape,kept),fill,workers,
413 tensor,permuted(Tensors.view_of(shape),List.append(&2,Nat,kept,reduced)),List.length(&2,Nat,reduced))
414 case count length: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},extents_at(shape,kept),fill,workers)
415
416# The reduction of the axes reduced, the others kept, both in increasing
417# index. A failed tensor keeps its problem; an invalid list of axes fails.
418def reduce_listed(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,name: String,
419 +operations: F32,+initial: Accumulator,tensor: Tensors.Tensor<Element>,+kept: List<&2,Nat>,+reduced: List<&2,Nat>,valid: Bool) -> Tensors.Tensor<Element>:
420 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
421 match status valid:
422 case Tensors.Usable{} True{}:
423 reduce_counted(~Element,~Accumulator,~operation,~finish,name,operations,Tensors.count_of(extents_at(shape,kept)),Tensors.count_of(extents_at(shape,reduced)),
424 initial,shape,kept,reduced,fill,workers,Tensors.Tensor{bands,geometry,shape,fill,Tensors.Usable{},workers})
425 case Tensors.Usable{} False{}: Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.InvalidArgument{shape}}},workers}
426 case Tensors.Failed{problem} valid: Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{problem},workers}
427
428# The result has the extents of the kept axes. Every listed axis is below the
429# rank and listed once. operations prices the combination of one value.
430def reduce_axes(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,name: String,
431 +operations: F32,+initial: Accumulator,tensor: Tensors.Tensor<Element>,+axes: List<&2,U32>) -> Tensors.Tensor<Element>:
432 Tensors.with_metadata(~Element,~Tensors.Tensor<Element>,tensor,tensor => shape => fill => status => workers =>
433 +rank = List.length(&2,U32,shape)
434 reduce_listed(~Element,~Accumulator,~operation,~finish,name,operations,initial,tensor,
435 axes_where(rank,0n,naturals(axes),False{}),axes_where(rank,0n,naturals(axes),True{}),
436 Views.unique_axes(naturals(axes)) && all_below(axes,rank)))
437
438# The axes index, index + 1, ... of a shape, remaining of them.
439def every_index(remaining: Nat,+index: Nat) -> List<&2,Nat>:
440 match remaining:
441 case 0n: []
442 case 1n+rest: index <> every_index(rest,1n+index)
443
444# The reduction of every element: no axis kept, every axis reduced.
445def reduce_all(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,name: String,
446 +operations: F32,+initial: Accumulator,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
447 Tensors.with_metadata(~Element,~Tensors.Tensor<Element>,tensor,tensor => shape => fill => status => workers =>
448 reduce_listed(~Element,~Accumulator,~operation,~finish,name,operations,initial,tensor,[],every_index(List.length(&2,U32,shape),0n),True{}))
449
450# Reversal along axes: one gather of the view of the tensor with the listed
451# axes read in reverse.
452def reversed_flags(remaining: Nat,+index: U32,+axes: List<&2,U32>) -> List<&2,Bool>:
453 match remaining:
454 case 0n: []
455 case 1n+rest: Views.contains_index(naturals(axes),U32.to_nat(index)) <> reversed_flags(rest,(index + 1 : U32),axes)
456
457def flipped(valid: Bool,view: Maybe<&2,Views.View>) -> Maybe<&2,Views.View>:
458 match valid:
459 case True{}: view
460 case False{}: None{}
461
462# The elements in reverse order along every listed axis; an axis is listed once.
463def flip(~Element: Data,tensor: Tensors.Tensor<Element>,+axes: List<&2,U32>) -> Tensors.Tensor<Element>:
464 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
465 +rank = List.length(&2,U32,shape)
466 Tensors.gather_extents(~Element,"flip",status,flipped(Views.unique_axes(naturals(axes)) && all_below(axes,rank),Tensors.view_of(shape)),
467 reversed_flags(rank,0,axes),bands,geometry,shape,fill,workers)
468
469# Padding. Constant fills the new elements with one value and writes the
470# tensor once, whatever the axes; Edge repeats the nearest element of the
471# tensor one axis after another, so a corner holds the nearest corner element.
472type Border<-Element: Data> is Data:
473 Constant{value: Element}
474 Edge{}
475
476def border_name(~Element: Data,border: Border<Element>) -> String:
477 match border:
478 case Constant{value}: "pad"
479 case Edge{}: "pad_edge"
480
481# The elements at one end of axis, repeated count times: (strip, tensor).
482def strip_of(~Element: Data,+axis: U32,from_end: Bool,+count: U32,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
483 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
484 +strip = with_extent(shape,U32.to_nat(axis),count)
485 Tensors.gather(~Element,"pad_edge",status,
486 Tensors.stretched(sliced(Tensors.view_of(shape),axis,Bool.pick(U32,from_end,(extent_at(shape,U32.to_nat(axis)) - 1 : U32),0),1,1),strip),[],
487 bands,geometry,shape,strip,fill,workers)
488
489# The elements at one end of axis, repeated count times: (strip, tensor).
490def edge_strip(~Element: Data,tensor: Tensors.Tensor<Element>,+axis: U32,from_end: Bool,+count: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
491 Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~(Tensors.Tensor<Element> & Tensors.Tensor<Element>),copy(~Element,tensor),original => copied =>
492 (strip_of(~Element,axis,from_end,count,copied),original))
493
494def pad_edge_axis(~Element: Data,unchanged: Bool,tensor: Tensors.Tensor<Element>,+axis: U32,+before: U32,+after: U32) -> Tensors.Tensor<Element>:
495 match unchanged:
496 case True{}: tensor
497 case False{}:
498 Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,edge_strip(~Element,tensor,axis,False{},before),first => tensor =>
499 Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,edge_strip(~Element,tensor,axis,True{},after),last => tensor =>
500 concat(~Element,Tensors.Cons{first,Tensors.Cons{tensor,Tensors.Single{last}}},axis)))
501
502def pad_edge_axes(~Element: Data,before: List<&2,U32>,after: List<&2,U32>,tensor: Tensors.Tensor<Element>,+axis: U32) -> Tensors.Tensor<Element>:
503 match before after:
504 case Con{+front,before} Con{+back,after}:
505 pad_edge_axes(~Element,before,after,pad_edge_axis(~Element,U32.is_zero(front) && U32.is_zero(back),tensor,axis,front,back),(axis + 1 : U32))
506 case before after: tensor
507
508# The view of the padded extents that holds the tensor: on every axis, the
509# extent of the tensor starting after the elements added in front.
510def interior(before: List<&2,U32>,shape: List<&2,U32>,view: Maybe<&2,Views.View>,+axis: U32) -> Maybe<&2,Views.View>:
511 match before shape:
512 case Con{front,before} Con{extent,shape}: interior(before,shape,sliced(view,axis,front,extent,1),(axis + 1 : U32))
513 case before shape: view
514
515def all_zero(values: List<&2,U32>) -> Bool:
516 match values:
517 case []: True{}
518 case value <> rest: U32.is_zero(value) && all_zero(rest)
519
520# Constant padding: storage of the padded extents filled with the value, then
521# the tensor written through the interior view. One copy whatever the axes.
522def padded_into(~Element: Data,count: Maybe<&2,U32>,+value: Element,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+padded: List<&2,U32>) -> Tensors.Tensor<Element>:
523 match count:
524 case Some{+count}:
525 Tensors.Tensor{part_bands,part_geometry,+part_shape,+part_fill,part_status,+part_workers} = tensor
526 +whole = {Bands.Geometry{1,1,1,count} : Bands.Geometry}
527 place_status(~Element,"pad",Tensors.Usable{},part_status,interior(before,part_shape,Tensors.view_of(padded),0),
528 Partition.Leaf{Tensors.buffer(~Element,value,count),0,1,Unit{}},whole,padded,part_shape,part_fill,part_workers,
529 Tensors.Tensor{part_bands,part_geometry,part_shape,part_fill,Tensors.Usable{},part_workers})
530 case None{}: Tensors.with_status(~Element,Tensors.Failed{Tensors.Problem{"pad",Tensors.InvalidShape{padded}}},tensor)
531
532def pad_constant(~Element: Data,unchanged: Bool,+value: Element,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+padded: List<&2,U32>) -> Tensors.Tensor<Element>:
533 match unchanged:
534 case True{}: tensor
535 case False{}: padded_into(~Element,Tensors.count_of(padded),value,tensor,before,padded)
536
537def pad_border(~Element: Data,border: Border<Element>,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor<Element>:
538 match border padded:
539 case Constant{value} Some{padded}: pad_constant(~Element,all_zero(before) && all_zero(after),value,tensor,before,padded)
540 case border padded: pad_edge_axes(~Element,before,after,tensor,0)
541
542def padded_count(shape: Maybe<&2,List<&2,U32>>) -> Maybe<&2,U32>:
543 match shape:
544 case Some{shape}: Tensors.count_of(shape)
545 case None{}: None{}
546
547def pad_valid(~Element: Data,+border: Border<Element>,status: Tensors.Status,valid: Bool,tensor: Tensors.Tensor<Element>,shape: List<&2,U32>,
548 before: List<&2,U32>,after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor<Element>:
549 match status valid:
550 case Tensors.Usable{} True{}: pad_border(~Element,border,tensor,before,after,padded)
551 case Tensors.Usable{} False{}: Tensors.with_status(~Element,Tensors.Failed{Tensors.Problem{border_name(~Element,border),Tensors.InvalidArgument{shape}}},tensor)
552 case Tensors.Failed{problem} valid: tensor
553
554# before[i] elements are added in front of axis i and after[i] behind it. Both
555# lists have one entry per axis and the padded extents must be allocatable.
556def pad(~Element: Data,+border: Border<Element>,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+after: List<&2,U32>) -> Tensors.Tensor<Element>:
557 Tensors.with_metadata(~Element,~Tensors.Tensor<Element>,tensor,tensor => shape => fill => status => workers =>
558 +shape = shape
559 +padded = Shapes.padded(shape,before,after)
560 pad_valid(~Element,border,status,Views.has_count(padded_count(padded)),tensor,shape,before,after,padded))