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. 5 import Base 6 import ./storage_buffer.bend as Storage 7 import ./tensor_shape.bend as Shapes 8 import ./tensor_view.bend as Views 9 import ./tensor_operations.bend as Operations 10 import ./tensor.bend as Tensors 11 import ./tensor_band.bend as Bands 12 import ./traversal_partition.bend as Partition 13 import ./tensor_band_concat.bend as BandConcat 14 import ./tensor_band_schedule.bend as BandSchedule 15 import ./tensor_band_view.bend as BandViews 16 import ./tensor_band_reduce.bend as BandReduce 17 18 def 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 23 def 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). 29 def 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 & Tensors.Tensor: 31 match status: 32 case Tensors.Usable{}: 33 Storage.bind_pair(~Bands.Bands(~Element),~Bands.Bands(~Element),~(Tensors.Tensor & Tensors.Tensor), 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 40 def copy(~Element: Data,tensor: Tensors.Tensor) -> Tensors.Tensor & Tensors.Tensor: 41 Tensors.Tensor{bands,geometry,shape,fill,status,workers} = tensor 42 copy_bands(~Element,status,bands,geometry,shape,fill,workers) 43 44 def 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]. 50 def reordered(~Element: Data,name: String,tensor: Tensors.Tensor,order: List<&2,Nat>) -> Tensors.Tensor: 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 54 def transpose(~Element: Data,tensor: Tensors.Tensor,order: List<&2,U32>) -> Tensors.Tensor: 55 reordered(~Element,"transpose",tensor,naturals(order)) 56 57 def 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. 63 def slice(~Element: Data,tensor: Tensors.Tensor,axis: U32,start: U32,count: U32,step: U32) -> Tensors.Tensor: 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 67 def 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. 74 def 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: 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 80 def 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: 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. 87 def reshape(~Element: Data,tensor: Tensors.Tensor,+shape: List<&2,U32>) -> Tensors.Tensor: 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 91 def extent_at(extents: List<&2,U32>,axis: Nat) -> U32: 92 Tensors.known_count(List.get(&2,U32,extents,axis)) 93 94 def 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. 100 def 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. 105 def 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 110 def 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 118 def 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 123 def 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 128 def 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 134 def 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: 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. 142 def 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 148 def 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: 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. 156 def upsample_nearest(~Element: Data,tensor: Tensors.Tensor,+scale: U32) -> Tensors.Tensor: 157 Tensors.Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor 158 upsample_status(~Element,status,bands,geometry,shape,scale,fill,workers) 159 160 def settle(~Element: Data,valid: Bool,result: Operations.Outcome<(Storage.Buffer & Storage.Buffer)>) -> 161 (Storage.Buffer & Storage.Buffer) & 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. 168 type Placing is Data: 169 Placing{banded: BandViews.Banded,geometry: Bands.Geometry} 170 171 def placed_outcome(~Element: Data,+first: U32,+rows: U32,result: Operations.Outcome<(Storage.Buffer & Storage.Buffer)>) -> 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 176 def placed_viewed(~Element: Data,source: Maybe<&2,Views.View>,slice: Views.View,+first: U32,+rows: U32,target: Storage.Buffer,part: Storage.Buffer) -> 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. 184 def placed_band(~Element: Data,context: Placing,first: U32,rows: U32,target: Storage.Buffer,part: Storage.Buffer) -> 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 189 def 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,owners: Bands.Bands(~Element) & Bands.Plan) -> Tensors.Tensor: 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,~Storage.Buffer,target,Tensors.arrive(~Element,BandViews.output(banded,geometry),layout,part)))) 195 196 def 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: 198 Tensors.Tensor{bands,geometry,target_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.ShapeMismatch{target_shape,part_shape}}},workers} 199 200 def 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,target: Tensors.Tensor) -> Tensors.Tensor: 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. 207 def 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) -> Tensors.Tensor: 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. 214 def 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) -> Tensors.Tensor: 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. 223 def 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) -> Tensors.Tensor: 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. 233 def 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) -> Tensors.Tensor: 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. 242 def 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) -> Tensors.Tensor: 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 250 def 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) -> Tensors.Tensor: 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. 259 def set_slice(~Element: Data,target: Tensors.Tensor,+axis: U32,start: U32,part: Tensors.Tensor) -> Tensors.Tensor: 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. 267 type Survey is Data: 268 Survey{valid: Bool,total: U32,status: Tensors.Status,workers: U32} 269 270 def 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 279 def prepend_survey(~Element: Data,tensor: Tensors.Tensor,result: Tensors.Parts & Survey) -> Tensors.Parts & 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. 284 def survey(~Element: Data,parts: Tensors.Parts,+first: List<&2,U32>,+axis: Nat,summary: Survey) -> Tensors.Parts & Survey: 285 match parts: 286 case Tensors.Single{tensor}: 287 Tensors.with_metadata(~Element,~(Tensors.Parts & 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 & 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 299 def with_extent(extents: List<&2,U32>,axis: Nat,+extent: U32) -> List<&2,U32>: 300 List.set(&2,U32,extents,axis,extent) 301 302 def appended_shaped(~Element: Data,+axis: U32,+start: U32,output: Tensors.Tensor,shaped: Tensors.Tensor & List<&2,U32>) -> Tensors.Tensor & U32: 303 (part,+shape) = shaped 304 (set_slice(~Element,output,axis,start,part),(start + extent_at(shape,U32.to_nat(axis)) : U32)) 305 306 def append_at(~Element: Data,+axis: U32,+start: U32,output: Tensors.Tensor,part: Tensors.Tensor) -> Tensors.Tensor & U32: 307 appended_shaped(~Element,axis,start,output,Tensors.shape(~Element,part)) 308 309 def append_part(~Element: Data,+axis: U32,state: Tensors.Tensor & U32,part: Tensors.Tensor) -> Tensors.Tensor & U32: 310 (output,start) = state 311 append_at(~Element,axis,start,output,part) 312 313 def appended_all(~Element: Data,state: Tensors.Tensor & U32) -> Tensors.Tensor: 314 (output,start) = state 315 output 316 317 def append_parts(~Element: Data,parts: Tensors.Parts,+axis: U32,state: Tensors.Tensor & U32) -> Tensors.Tensor: 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. 324 def band_joined(~Element: Data,parts: Tensors.Parts,shape: List<&2,U32>,fill: Element) -> Tensors.Tensor: 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. 329 def join_counted(~Element: Data,count: Maybe<&2,U32>,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,parts: Tensors.Parts) -> Tensors.Tensor: 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. 337 def join_shaped(~Element: Data,empty: Bool,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,parts: Tensors.Parts) -> Tensors.Tensor: 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 342 def 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) -> Tensors.Tensor: 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 349 def joined(~Element: Data,+first: List<&2,U32>,+axis: U32,+fill: Element,surveyed: Tensors.Parts & Survey) -> Tensors.Tensor: 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 354 def lead(~Element: Data,parts: Tensors.Parts) -> Tensors.Parts & (List<&2,U32> & Element): 355 match parts: 356 case Tensors.Single{tensor}: 357 Tensors.with_metadata(~Element,~(Tensors.Parts & (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 & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers => 363 +shape = shape 364 +fill = fill 365 (Tensors.Cons{tensor,rest},(shape,fill))) 366 367 def concat_first(~Element: Data,+axis: U32,parts: Tensors.Parts,+shape: List<&2,U32>,+fill: Element) -> Tensors.Tensor: 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 371 def concat_led(~Element: Data,+axis: U32,led: Tensors.Parts & (List<&2,U32> & Element)) -> Tensors.Tensor: 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. 377 def concat(~Element: Data,parts: Tensors.Parts,+axis: U32) -> Tensors.Tensor: 378 concat_led(~Element,axis,lead(~Element,parts)) 379 380 # Every listed axis is below the rank. 381 def 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. 391 def 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 397 def 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 402 def 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 407 def 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) -> Tensors.Tensor: 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. 418 def reduce_listed(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,name: String, 419 +operations: F32,+initial: Accumulator,tensor: Tensors.Tensor,+kept: List<&2,Nat>,+reduced: List<&2,Nat>,valid: Bool) -> Tensors.Tensor: 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. 430 def reduce_axes(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,name: String, 431 +operations: F32,+initial: Accumulator,tensor: Tensors.Tensor,+axes: List<&2,U32>) -> Tensors.Tensor: 432 Tensors.with_metadata(~Element,~Tensors.Tensor,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. 439 def 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. 445 def reduce_all(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element,name: String, 446 +operations: F32,+initial: Accumulator,tensor: Tensors.Tensor) -> Tensors.Tensor: 447 Tensors.with_metadata(~Element,~Tensors.Tensor,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. 452 def 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 457 def 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. 463 def flip(~Element: Data,tensor: Tensors.Tensor,+axes: List<&2,U32>) -> Tensors.Tensor: 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. 472 type Border<-Element: Data> is Data: 473 Constant{value: Element} 474 Edge{} 475 476 def border_name(~Element: Data,border: Border) -> 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). 482 def strip_of(~Element: Data,+axis: U32,from_end: Bool,+count: U32,tensor: Tensors.Tensor) -> Tensors.Tensor: 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). 490 def edge_strip(~Element: Data,tensor: Tensors.Tensor,+axis: U32,from_end: Bool,+count: U32) -> Tensors.Tensor & Tensors.Tensor: 491 Tensors.then(~Tensors.Tensor,~Tensors.Tensor,~(Tensors.Tensor & Tensors.Tensor),copy(~Element,tensor),original => copied => 492 (strip_of(~Element,axis,from_end,count,copied),original)) 493 494 def pad_edge_axis(~Element: Data,unchanged: Bool,tensor: Tensors.Tensor,+axis: U32,+before: U32,+after: U32) -> Tensors.Tensor: 495 match unchanged: 496 case True{}: tensor 497 case False{}: 498 Tensors.then(~Tensors.Tensor,~Tensors.Tensor,~Tensors.Tensor,edge_strip(~Element,tensor,axis,False{},before),first => tensor => 499 Tensors.then(~Tensors.Tensor,~Tensors.Tensor,~Tensors.Tensor,edge_strip(~Element,tensor,axis,True{},after),last => tensor => 500 concat(~Element,Tensors.Cons{first,Tensors.Cons{tensor,Tensors.Single{last}}},axis))) 501 502 def pad_edge_axes(~Element: Data,before: List<&2,U32>,after: List<&2,U32>,tensor: Tensors.Tensor,+axis: U32) -> Tensors.Tensor: 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. 510 def 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 515 def 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. 522 def padded_into(~Element: Data,count: Maybe<&2,U32>,+value: Element,tensor: Tensors.Tensor,+before: List<&2,U32>,+padded: List<&2,U32>) -> Tensors.Tensor: 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 532 def pad_constant(~Element: Data,unchanged: Bool,+value: Element,tensor: Tensors.Tensor,+before: List<&2,U32>,+padded: List<&2,U32>) -> Tensors.Tensor: 533 match unchanged: 534 case True{}: tensor 535 case False{}: padded_into(~Element,Tensors.count_of(padded),value,tensor,before,padded) 536 537 def pad_border(~Element: Data,border: Border,tensor: Tensors.Tensor,+before: List<&2,U32>,+after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor: 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 542 def 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 547 def pad_valid(~Element: Data,+border: Border,status: Tensors.Status,valid: Bool,tensor: Tensors.Tensor,shape: List<&2,U32>, 548 before: List<&2,U32>,after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor: 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. 556 def pad(~Element: Data,+border: Border,tensor: Tensors.Tensor,+before: List<&2,U32>,+after: List<&2,U32>) -> Tensors.Tensor: 557 Tensors.with_metadata(~Element,~Tensors.Tensor,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))