tensor_layout.bend source

tensor_layout.bend on the hub · documented module

# Layout operations consume their arguments. Copy and nearest upsampling# preserve row bands. Other layout operations first obtain one dense owner;# reshape relabels it, while the remaining operations use checked view copies.import Baseimport ./storage_buffer.bend as Storageimport ./tensor_shape.bend as Shapesimport ./tensor_view.bend as Viewsimport ./tensor_operations.bend as Operationsimport ./tensor.bend as Tensorsimport ./tensor_band.bend as Bandsimport ./traversal_partition.bend as Partitionimport ./tensor_band_concat.bend as BandConcatimport ./tensor_band_schedule.bend as BandScheduledef extents(view: Maybe<&2,Views.View>) -> List<&2,U32>:  match view:    case Some{view}: Views.axis_extents_of(view)    case None{}: []def naturals(values: List<&2,U32>) -> List<&2,Nat>:  match values:    case []: []    case value <> rest: U32.to_nat(value) <> naturals(rest)def copied(~Element: Data,name: String,+source_shape: List<&2,U32>,+target_shape: List<&2,U32>,+fill: Element,+workers: U32,  result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match result:    case Operations.Accepted{(input,output)}:      (Tensors.Tensor{output,target_shape,fill,Tensors.Usable{},workers},Tensors.Tensor{input,source_shape,fill,Tensors.Usable{},workers})    case Operations.Rejected{(input,output)}:      (Tensors.Tensor{output,target_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.InvalidArgument{source_shape}}},workers},        Tensors.Tensor{input,source_shape,fill,Tensors.Usable{},workers})def gather_views(~Element: Data,name: String,source_view: Maybe<&2,Views.View>,target_view: Maybe<&2,Views.View>,count: Maybe<&2,U32>,  input: Storage.Buffer<Element>,+source_shape: List<&2,U32>,+target_shape: List<&2,U32>,+fill: Element,+workers: U32) ->  Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match source_view target_view count:    case Some{source_view} Some{target_view} Some{count}:      copied(~Element,name,source_shape,target_shape,fill,workers,Operations.copy(~Element,input,source_view,Tensors.buffer(~Element,fill,count),target_view))    case source_view target_view count:      (Tensors.failed(~Element,name,Tensors.InvalidArgument{source_shape},target_shape,fill,workers),Tensors.Tensor{input,source_shape,fill,Tensors.Usable{},workers})# Copy the elements selected by source_view, in its row-major order, into a# new tensor whose storage is laid out as view_shape and labelled target_shape.def gather_status(~Element: Data,name: String,status: Tensors.Status,source_view: Maybe<&2,Views.View>,+view_shape: List<&2,U32>,  +target_shape: List<&2,U32>,input: Storage.Buffer<Element>,+source_shape: List<&2,U32>,+fill: Element,+workers: U32) ->  Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match status:    case Tensors.Usable{}:      gather_views(~Element,name,source_view,Tensors.view_of(view_shape),Tensors.count_of(view_shape),input,source_shape,target_shape,fill,workers)    case Tensors.Failed{+problem}:      (Tensors.Tensor{Tensors.buffer(~Element,fill,0),target_shape,fill,Tensors.Failed{problem},workers},        Tensors.Tensor{input,source_shape,fill,Tensors.Failed{problem},workers})def result_of(~Element: Data,pair: Tensors.Tensor<Element> & Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:  (result,source) = pair  resultdef swapped(~Element: Data,pair: Tensors.Tensor<Element> & Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  (result,source) = pair  (source,result)# The only layout operation that keeps its argument: (original, copy).def copy_bands(~Element: Data,status: Tensors.Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+fill: Element,+workers: U32) ->  Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match status:    case Tensors.Usable{}:      Storage.bind_pair(~Bands.Bands(~Element),~Bands.Bands(~Element),~(Tensors.Tensor<Element> & Tensors.Tensor<Element>),        Bands.copy(~Element,bands),left => right =>          (Tensors.Banded{left,geometry,shape,fill,Tensors.Usable{},workers},Tensors.Banded{right,geometry,shape,fill,Tensors.Usable{},workers}))    case Tensors.Failed{+problem}:      (Tensors.Banded{bands,geometry,shape,fill,Tensors.Failed{problem},workers},        Tensors.Tensor{Tensors.buffer(~Element,fill,0),shape,fill,Tensors.Failed{problem},workers})def copy(~Element: Data,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match tensor:    case Tensors.Tensor{storage,+shape,fill,status,workers}:      swapped(~Element,gather_status(~Element,"copy",status,Tensors.view_of(shape),shape,shape,storage,shape,fill,workers))    case Tensors.Banded{bands,geometry,shape,fill,status,workers}: copy_bands(~Element,status,bands,geometry,shape,fill,workers)def permuted(view: Maybe<&2,Views.View>,order: List<&2,Nat>) -> Maybe<&2,Views.View>:  match view:    case Some{view}: Views.permute(view,order)    case None{}: None{}def transpose_view(~Element: Data,+source: Maybe<&2,Views.View>,status: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element,  +workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  gather_status(~Element,"transpose",status,source,extents(source),extents(source),storage,shape,fill,workers)# The result's axis i is the source axis order[i].def transpose(~Element: Data,tensor: Tensors.Tensor<Element>,order: List<&2,U32>) -> Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>    +shape = shape    result_of(~Element,transpose_view(~Element,permuted(Tensors.view_of(shape),naturals(order)),status,storage,shape,fill,workers)))def sliced(view: Maybe<&2,Views.View>,axis: U32,start: U32,count: U32,step: U32) -> Maybe<&2,Views.View>:  match view:    case Some{view}: Views.slice(view,U32.to_nat(axis),start,count,step)    case None{}: None{}def slice_view(~Element: Data,+source: Maybe<&2,Views.View>,status: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element,  +workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  gather_status(~Element,"slice",status,source,extents(source),extents(source),storage,shape,fill,workers)# Elements start, start + step, ... (count of them) along one axis.def slice(~Element: Data,tensor: Tensors.Tensor<Element>,axis: U32,start: U32,count: U32,step: U32) -> Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>    +shape = shape    result_of(~Element,slice_view(~Element,sliced(Tensors.view_of(shape),axis,start,count,step),status,storage,shape,fill,workers)))def same_count(left: Maybe<&2,U32>,right: Maybe<&2,U32>) -> Bool:  match left right:    case Some{left} Some{right}: U32.is_eq(left,right)    case left right: False{}def relabel_ready(~Element: Data,accepted: Bool,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+target: List<&2,U32>,fill: Element,  workers: U32) -> Tensors.Tensor<Element>:  match accepted:    case True{}: Tensors.Tensor{storage,target,fill,Tensors.Usable{},workers}    case False{}: Tensors.Tensor{storage,shape,fill,Tensors.Failed{Tensors.Problem{"reshape",Tensors.ShapeMismatch{shape,target}}},workers}def relabel(~Element: Data,accepted: Bool,current: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+target: List<&2,U32>,fill: Element,  workers: U32) -> Tensors.Tensor<Element>:  match current:    case Tensors.Usable{}: relabel_ready(~Element,accepted,storage,shape,target,fill,workers)    case Tensors.Failed{problem}: Tensors.Tensor{storage,shape,fill,Tensors.Failed{problem},workers}# Same elements in the same row-major order under new extents.def reshape(~Element: Data,tensor: Tensors.Tensor<Element>,+shape: List<&2,U32>) -> Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => source_shape => fill => status => workers =>    +source_shape = source_shape    relabel(~Element,same_count(Tensors.count_of(source_shape),Tensors.count_of(shape)),status,storage,source_shape,shape,fill,workers))# Nearest-neighbour upsampling repeats each element of the last two axes# scale times, by copying from a view whose repeated axes have stride zero.def prepend_repeated(axis: Views.Axis,rest: Maybe<&2,List<&2,Views.Axis>>) -> Maybe<&2,List<&2,Views.Axis>>:  match rest:    case Some{axes}: Some{axis <> axes}    case None{}: None{}def repeated_axes(axes: List<&2,Views.Axis>,+scale: U32) -> Maybe<&2,List<&2,Views.Axis>>:  match axes:    case [Views.Axis{height,height_stride},Views.Axis{width,width_stride}]:      Some{[Views.Axis{height,height_stride},Views.Axis{scale,0},Views.Axis{width,width_stride},Views.Axis{scale,0}]}    case [axis] : None{}    case []: None{}    case axis <> rest: prepend_repeated(axis,repeated_axes(rest,scale))def with_offset(axes: Maybe<&2,List<&2,Views.Axis>>,offset: U32) -> Maybe<&2,Views.View>:  match axes:    case Some{axes}: Some{Views.View{axes,offset}}    case None{}: None{}def repeated(view: Maybe<&2,Views.View>,scale: U32) -> Maybe<&2,Views.View>:  match view:    case Some{Views.View{axes,offset}}: with_offset(repeated_axes(axes,scale),offset)    case None{}: None{}def scaled_extents(extents: List<&2,U32>,+scale: U32) -> List<&2,U32>:  match extents:    case [height,width]: [(height * scale : U32),(width * scale : U32)]    case extent <> rest: extent <> scaled_extents(rest,scale)    case []: []def upsample_view(~Element: Data,+source: Maybe<&2,Views.View>,status: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+scale: U32,fill: Element,  +workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  gather_status(~Element,"upsample_nearest",status,source,extents(source),scaled_extents(shape,scale),storage,shape,fill,workers)type Upsample<-Element: Data> is Data:  Upsample{geometry: Bands.Geometry,scale: U32,fill: Element}def upsample_leaf(~Element: Data,context: Upsample<Element>,first: U32,rows: U32,input: Storage.Buffer<Element>) -> Bands.Bands(~Element):  Upsample{Bands.Geometry{batch,channels,height,width},+scale,fill} = context  +rows = rows  +shape = {[channels,rows,width] : List<&2,U32>}  +first = (first * scale : U32)  +rows = (rows * scale : U32)  Tensors.with_dense(~Element,~Bands.Bands(~Element),    result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),Tensors.Usable{},input,shape,scale,fill,1)),    output => shape => fill => status => workers => Partition.Leaf{output,first,rows,Unit{}})def upsample_supported(~Element: Data,valid: Bool,bands: Bands.Bands(~Element),geometry: Bands.Geometry,+shape: List<&2,U32>,+scale: U32,  +fill: Element,+workers: U32) -> Tensors.Tensor<Element>:  match valid geometry:    case False{} geometry:      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,Tensors.Banded{bands,geometry,shape,fill,Tensors.Usable{},workers},        storage => shape => fill => status => workers =>          +shape = shape          result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),status,storage,shape,scale,fill,workers)))    case True{} Bands.Geometry{+batch,+channels,+height,+width}:      Tensors.Banded{BandSchedule.transform(~Element,~Upsample<Element>,~upsample_leaf(~Element),Upsample{Bands.Geometry{batch,channels,height,width},scale,fill},          Bands.Geometry{batch,channels,height,width},workers,(U32.to_f32(scale) * U32.to_f32(scale) * 12.0 : F32),bands),        Bands.Geometry{batch,channels,(height * scale : U32),(width * scale : U32)},scaled_extents(shape,scale),fill,Tensors.Usable{},workers}def upsample_bands(~Element: Data,status: Tensors.Status,bands: Bands.Bands(~Element),geometry: Bands.Geometry,+shape: List<&2,U32>,+scale: U32,  +fill: Element,workers: U32) -> Tensors.Tensor<Element>:  match status geometry:    case Tensors.Failed{+problem} geometry:      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,Tensors.Banded{bands,geometry,shape,fill,Tensors.Failed{problem},workers},        storage => shape => fill => status => workers =>          +shape = shape          result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),status,storage,shape,scale,fill,workers)))    case Tensors.Usable{} Bands.Geometry{batch,channels,+height,+width}:      upsample_supported(~Element,Bool.not(U32.is_zero(scale)) && Views.has_count(Shapes.checked_multiply(height,scale)) &&        Views.has_count(Shapes.checked_multiply(width,scale)) && Views.has_count(Tensors.count_of(scaled_extents(shape,scale))) &&        Bool.not(Shapes.has_zero(scaled_extents(shape,scale))),        bands,Bands.Geometry{batch,channels,height,width},shape,scale,fill,workers)def upsample_nearest(~Element: Data,tensor: Tensors.Tensor<Element>,+scale: U32) -> Tensors.Tensor<Element>:  match tensor:    case Tensors.Tensor{storage,+shape,fill,status,workers}:      result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),status,storage,shape,scale,fill,workers))    case Tensors.Banded{bands,geometry,shape,fill,status,workers}: upsample_bands(~Element,status,bands,geometry,shape,scale,fill,workers)def settle(~Element: Data,valid: Bool,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) ->  (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:  match result:    case Operations.Accepted{owners}: (owners,valid)    case Operations.Rejected{owners}: (owners,False{})def count_or_zero(count: Maybe<&2,U32>) -> U32:  match count:    case Some{count}: count    case None{}: 0def extent_at(extents: List<&2,U32>,axis: Nat) -> U32:  count_or_zero(List.get(&2,U32,extents,axis))def settled(~Element: Data,name: String,valid: Bool,part: Storage.Buffer<Element>,target: Storage.Buffer<Element>,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,  target_fill: Element,part_fill: Element,+workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match valid:    case True{}: (Tensors.Tensor{target,target_shape,target_fill,Tensors.Usable{},workers},Tensors.Tensor{part,part_shape,part_fill,Tensors.Usable{},workers})    case False{}:      (Tensors.Tensor{target,target_shape,target_fill,Tensors.Failed{Tensors.Problem{name,Tensors.ShapeMismatch{target_shape,part_shape}}},workers},        Tensors.Tensor{part,part_shape,part_fill,Tensors.Usable{},workers})def placed(~Element: Data,name: String,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,target_fill: Element,part_fill: Element,+workers: U32,  result: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  (owners,valid) = result  (part,target) = owners  settled(~Element,name,valid,part,target,target_shape,part_shape,target_fill,part_fill,workers)def same_extents(slice: Views.View,part_shape: List<&2,U32>) -> Bool:  Shapes.same(Views.shape(slice),Shapes.Shape{part_shape})# The part written through a view of the target storage: one checked copy,# which rejects a view whose extents differ from the part. A part with no# element writes nothing; its extents are compared here.def place(~Element: Data,name: String,empty: Bool,slice: Maybe<&2,Views.View>,source: Maybe<&2,Views.View>,  target: Storage.Buffer<Element>,+target_shape: List<&2,U32>,target_fill: Element,  part: Storage.Buffer<Element>,+part_shape: List<&2,U32>,part_fill: Element,+workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match empty slice source:    case False{} Some{slice} Some{source}:      placed(~Element,name,target_shape,part_shape,target_fill,part_fill,workers,settle(~Element,True{},Operations.scatter(~Element,part,source,target,slice)))    case True{} Some{slice} source:      settled(~Element,name,same_extents(slice,part_shape),part,target,target_shape,part_shape,target_fill,part_fill,workers)    case empty slice source: settled(~Element,name,False{},part,target,target_shape,part_shape,target_fill,part_fill,workers)def place_status(~Element: Data,name: String,target_status: Tensors.Status,part_status: Tensors.Status,slice: Maybe<&2,Views.View>,  target: Storage.Buffer<Element>,+target_shape: List<&2,U32>,+target_fill: Element,  part: Storage.Buffer<Element>,+part_shape: List<&2,U32>,+part_fill: Element,+workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  match target_status part_status:    case Tensors.Usable{} Tensors.Usable{}:      place(~Element,name,Shapes.has_zero(part_shape),slice,Tensors.view_of(part_shape),target,target_shape,target_fill,part,part_shape,part_fill,workers)    case Tensors.Usable{} Tensors.Failed{+problem}:      (Tensors.Tensor{target,target_shape,target_fill,Tensors.Failed{problem},workers},Tensors.Tensor{part,part_shape,part_fill,Tensors.Failed{problem},workers})    case Tensors.Failed{problem} part_status:      (Tensors.Tensor{target,target_shape,target_fill,Tensors.Failed{problem},workers},Tensors.Tensor{part,part_shape,part_fill,part_status,workers})# target[..., start : start + part extent, ...] = part along axis: the part is# written through the slice of the target that it replaces.def set_slice(~Element: Data,target: Tensors.Tensor<Element>,+axis: U32,start: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,target,target => target_shape => target_fill => target_status => target_workers =>    +target_shape = target_shape    Tensors.with_dense(~Element,~Tensors.Tensor<Element>,part,part => part_shape => part_fill => part_status => part_workers =>      +part_shape = part_shape      result_of(~Element,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),target,target_shape,target_fill,part,part_shape,part_fill,        U32.max(target_workers,part_workers)))))# Concatenation along an axis.type Survey is Data:  Survey{valid: Bool,total: U32,status: Tensors.Status,workers: U32}def surveyed(survey: Survey,+first: List<&2,U32>,+shape: List<&2,U32>,+axis: Nat,status: Tensors.Status,workers: U32) -> Survey:  Survey{valid,+total,previous,most} = survey  +width = extent_at(shape,axis)  Survey{valid && Nat.is_eq(List.length(&2,U32,first),List.length(&2,U32,shape)) &&      Shapes.same_extents(List.take(&2,U32,first,axis),List.take(&2,U32,shape,axis)) &&      Shapes.same_extents(List.drop(&2,U32,first,1n+axis),List.drop(&2,U32,shape,1n+axis)) &&      U32.is_le(width,(4294967295 - total : U32)),    (total + width : U32),Tensors.first_failure(previous,status),U32.max(most,workers)}def prepend_survey(~Element: Data,tensor: Tensors.Tensor<Element>,result: Tensors.Parts<Element> & Survey) -> Tensors.Parts<Element> & Survey:  (rest,summary) = result  (Tensors.Cons{tensor,rest},summary)# Summarizes the extents of every part and returns the parts in order.def survey(~Element: Data,parts: Tensors.Parts<Element>,+first: List<&2,U32>,+axis: Nat,summary: Survey) -> Tensors.Parts<Element> & Survey:  match parts:    case Tensors.Single{tensor}:      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & Survey),tensor,tensor => shape => fill => status => workers =>        +shape = shape        +status = status        +workers = workers        (Tensors.Single{tensor},surveyed(summary,first,shape,axis,status,workers)))    case Tensors.Cons{tensor,rest}:      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & Survey),tensor,tensor => shape => fill => status => workers =>        +shape = shape        +status = status        +workers = workers        prepend_survey(~Element,tensor,survey(~Element,rest,first,axis,surveyed(summary,first,shape,axis,status,workers))))def with_extent(extents: List<&2,U32>,axis: Nat,+extent: U32) -> List<&2,U32>:  List.set(&2,U32,extents,axis,extent)def append_at(~Element: Data,+axis: U32,+start: U32,output: Tensors.Tensor<Element>,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & U32:  Tensors.with_dense(~Element,~(Tensors.Tensor<Element> & U32),part,storage => shape => fill => status => workers =>    +shape = shape    (set_slice(~Element,output,axis,start,Tensors.Tensor{storage,shape,fill,status,workers}),(start + extent_at(shape,U32.to_nat(axis)) : U32)))def append_part(~Element: Data,+axis: U32,state: Tensors.Tensor<Element> & U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & U32:  (output,start) = state  append_at(~Element,axis,start,output,part)def appended_all(~Element: Data,state: Tensors.Tensor<Element> & U32) -> Tensors.Tensor<Element>:  (output,start) = state  outputdef append_parts(~Element: Data,parts: Tensors.Parts<Element>,+axis: U32,state: Tensors.Tensor<Element> & U32) -> Tensors.Tensor<Element>:  match parts:    case Tensors.Single{part}: appended_all(~Element,append_part(~Element,axis,state,part))    case Tensors.Cons{part,rest}: append_parts(~Element,rest,axis,append_part(~Element,axis,state,part))def join_selected(~Element: Data,geometry: Maybe<&2,Bands.Geometry>,count: Maybe<&2,U32>,shape: List<&2,U32>,axis: U32,fill: Element,workers: U32,  layout: Maybe<&2,Bands.Plan>,parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:  match geometry count layout:    case Some{geometry} Some{count} Some{layout}: BandConcat.join(~Element,layout,geometry,parts,shape,fill,workers)    case geometry count layout: append_parts(~Element,parts,axis,(Tensors.with_workers(~Element,Tensors.full(~Element,shape,fill),workers),0))def join_discovered(~Element: Data,+shape: List<&2,U32>,+axis: U32,fill: Element,workers: U32,  discovered: Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) -> Tensors.Tensor<Element>:  (parts,layout) = discovered  join_selected(~Element,BandConcat.geometry(shape,axis),Tensors.count_of(shape),shape,axis,fill,workers,layout,parts)# The final branch is one Bool. Plan comparison stays in the metadata pass,# outside every element-copy loop. Dense append gathers each band source once.def join_compatible_selected(~Element: Data,compatible: Bool,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,  discovered: Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) -> Tensors.Tensor<Element>:  match compatible:    case True{}: join_discovered(~Element,shape,axis,fill,workers,discovered)    case False{}:      (parts,layout) = discovered      append_parts(~Element,parts,axis,(Tensors.with_workers(~Element,Tensors.full(~Element,shape,fill),workers),0))def join_compatible(~Element: Data,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,  discovered: (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool) -> Tensors.Tensor<Element>:  (parts_and_layout,compatible) = discovered  join_compatible_selected(~Element,compatible && Bool.not(Shapes.has_zero(shape)),shape,axis,fill,workers,parts_and_layout)def join_status(~Element: Data,status: Tensors.Status,valid: Bool,+first: List<&2,U32>,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,  parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:  match status valid:    case Tensors.Usable{} True{}: join_compatible(~Element,shape,axis,fill,workers,BandConcat.discover_compatible(~Element,parts))    case Tensors.Usable{} False{}: Tensors.failed(~Element,"concat",Tensors.InvalidArgument{first},first,fill,workers)    case Tensors.Failed{problem} other: Tensors.Tensor{Tensors.buffer(~Element,fill,0),first,fill,Tensors.Failed{problem},workers}def joined(~Element: Data,+first: List<&2,U32>,+axis: U32,+fill: Element,surveyed: Tensors.Parts<Element> & Survey) -> Tensors.Tensor<Element>:  (parts,summary) = surveyed  Survey{valid,total,status,workers} = summary  join_status(~Element,status,valid,first,with_extent(first,U32.to_nat(axis),total),axis,fill,workers,parts)def lead(~Element: Data,parts: Tensors.Parts<Element>) -> Tensors.Parts<Element> & (List<&2,U32> & Element):  match parts:    case Tensors.Single{tensor}:      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers =>        +shape = shape        +fill = fill        (Tensors.Single{tensor},(shape,fill)))    case Tensors.Cons{tensor,rest}:      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers =>        +shape = shape        +fill = fill        (Tensors.Cons{tensor,rest},(shape,fill)))def concat_first(~Element: Data,+axis: U32,parts: Tensors.Parts<Element>,+shape: List<&2,U32>,+fill: Element) -> Tensors.Tensor<Element>:  joined(~Element,shape,axis,fill,    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}))def concat_led(~Element: Data,+axis: U32,led: Tensors.Parts<Element> & (List<&2,U32> & Element)) -> Tensors.Tensor<Element>:  (parts,first) = led  (shape,fill) = first  concat_first(~Element,axis,parts,shape,fill)# Parts agree on every extent except axis; the result stacks them in order.def concat(~Element: Data,parts: Tensors.Parts<Element>,+axis: U32) -> Tensors.Tensor<Element>:  concat_led(~Element,axis,lead(~Element,parts))# Broadcasting to given extents: one checked copy from a stretched view.def broadcast_status(~Element: Data,status: Tensors.Status,view: Maybe<&2,Views.View>,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,  +target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:  match status view:    case Tensors.Usable{} None{}: Tensors.failed(~Element,"broadcast_to",Tensors.ShapeMismatch{shape,target},target,fill,workers)    case status view: result_of(~Element,gather_status(~Element,"broadcast_to",status,view,target,target,storage,shape,fill,workers))# The tensor repeated to the extents target under NumPy broadcasting.def broadcast_to(~Element: Data,tensor: Tensors.Tensor<Element>,+target: List<&2,U32>) -> Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>    +shape = shape    broadcast_status(~Element,status,Tensors.stretched(Tensors.view_of(shape),target),storage,shape,target,fill,workers))# Reduction along one axis. The output starts at the initial value and takes# the slices of the axis in increasing index, so each output element is# operation(... operation(operation(initial, x[0]), x[1]) ..., x[extent - 1]).def fold_slice(~Element: Data,~operation: Element -> Element -> Element,valid: Bool,slice: Maybe<&2,Views.View>,target: Maybe<&2,Views.View>,  output: Storage.Buffer<Element>,input: Storage.Buffer<Element>) -> (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:  match slice target:    case Some{slice} Some{target}: settle(~Element,valid,Operations.combine(~Element,~operation,output,target,input,slice))    case slice target: ((output,input),False{})# The views are built from the extents at every step: a loop that carried one# View value through its iterations would make the compiler share every View# constructor of the program.def fold_step(~Element: Data,~operation: Element -> Element -> Element,index: U32,axis: U32,shape: List<&2,U32>,unit: List<&2,U32>,  state: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:  (owners,valid) = state  (output,input) = owners  fold_slice(~Element,~operation,valid,sliced(Tensors.view_of(shape),axis,index,1,1),Tensors.view_of(unit),output,input)def fold_slices(~Element: Data,~operation: Element -> Element -> Element,remaining: Nat,+index: U32,+axis: U32,+shape: List<&2,U32>,+unit: List<&2,U32>,  state: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:  match remaining:    case 0n: state    case 1n+rest:      fold_slices(~Element,~operation,rest,(index + 1 : U32),axis,shape,unit,fold_step(~Element,~operation,index,axis,shape,unit,state))def folded_valid(~Element: Data,name: String,valid: Bool,output: Storage.Buffer<Element>,shape: List<&2,U32>,result_shape: List<&2,U32>,fill: Element,  workers: U32) -> Tensors.Tensor<Element>:  match valid:    case True{}: Tensors.Tensor{output,result_shape,fill,Tensors.Usable{},workers}    case False{}: Tensors.Tensor{output,result_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.InvalidArgument{shape}}},workers}def folded(~Element: Data,name: String,shape: List<&2,U32>,result_shape: List<&2,U32>,fill: Element,workers: U32,  result: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> Tensors.Tensor<Element>:  (owners,valid) = result  (output,input) = owners  folded_valid(~Element,name,valid,output,shape,result_shape,fill,workers)# unit is the shape with extent 1 at the axis: the layout of the output and# the extents of every slice.def reduce_axis_count(~Element: Data,~operation: Element -> Element -> Element,name: String,count: Maybe<&2,U32>,+axis: U32,input: Storage.Buffer<Element>,  +shape: List<&2,U32>,unit: List<&2,U32>,+initial: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:  match count:    case Some{count}:      folded(~Element,name,shape,Shapes.without_axis(shape,U32.to_nat(axis)),fill,workers,        fold_slices(~Element,~operation,U32.to_nat(extent_at(shape,U32.to_nat(axis))),0,axis,shape,unit,((Tensors.buffer(~Element,initial,count),input),True{})))    case None{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},Shapes.without_axis(shape,U32.to_nat(axis)),fill,workers)def reduce_axis_status(~Element: Data,~operation: Element -> Element -> Element,name: String,status: Tensors.Status,valid: Bool,+axis: U32,  input: Storage.Buffer<Element>,+shape: List<&2,U32>,+initial: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:  match status valid:    case Tensors.Usable{} True{}:      reduce_axis_count(~Element,~operation,name,Tensors.count_of(with_extent(shape,U32.to_nat(axis),1)),axis,input,shape,with_extent(shape,U32.to_nat(axis),1),        initial,fill,workers)    case Tensors.Usable{} False{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},shape,fill,workers)    case Tensors.Failed{problem} valid: Tensors.Tensor{input,shape,fill,Tensors.Failed{problem},workers}# The result has the extents of the tensor without axis.def reduce_axis(~Element: Data,~operation: Element -> Element -> Element,name: String,+initial: Element,tensor: Tensors.Tensor<Element>,+axis: U32) ->  Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>    +shape = shape    reduce_axis_status(~Element,~operation,name,status,Nat.is_lt(U32.to_nat(axis),List.length(&2,U32,shape)),axis,storage,shape,initial,fill,workers))# Padding. Constant fills the new elements with one value and writes the# tensor once, whatever the axes; Edge repeats the nearest element of the# tensor one axis after another, so a corner holds the nearest corner element.type Border<-Element: Data> is Data:  Constant{value: Element}  Edge{}def border_name(~Element: Data,border: Border<Element>) -> String:  match border:    case Constant{value}: "pad"    case Edge{}: "pad_edge"def blank(~Element: Data,count: Maybe<&2,U32>,+shape: List<&2,U32>,+value: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:  match count:    case Some{count}: Tensors.Tensor{Tensors.buffer(~Element,value,count),shape,fill,Tensors.Usable{},workers}    case None{}: Tensors.failed(~Element,"pad",Tensors.InvalidShape{shape},shape,fill,workers)# The elements at one end of axis, repeated count times: (strip, tensor).def edge_strip(~Element: Data,tensor: Tensors.Tensor<Element>,+axis: U32,from_end: Bool,+count: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:  Tensors.with_dense(~Element,~(Tensors.Tensor<Element> & Tensors.Tensor<Element>),tensor,storage => shape => fill => status => workers =>    +shape = shape    +strip = with_extent(shape,U32.to_nat(axis),count)    gather_status(~Element,"pad_edge",status,      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),      strip,strip,storage,shape,fill,workers))def pad_edge_axis(~Element: Data,unchanged: Bool,tensor: Tensors.Tensor<Element>,+axis: U32,+before: U32,+after: U32) -> Tensors.Tensor<Element>:  match unchanged:    case True{}: tensor    case False{}:      Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,edge_strip(~Element,tensor,axis,False{},before),first => tensor =>        Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,edge_strip(~Element,tensor,axis,True{},after),last => tensor =>          concat(~Element,Tensors.Cons{first,Tensors.Cons{tensor,Tensors.Single{last}}},axis)))def pad_edge_axes(~Element: Data,before: List<&2,U32>,after: List<&2,U32>,tensor: Tensors.Tensor<Element>,+axis: U32) -> Tensors.Tensor<Element>:  match before after:    case Con{+front,before} Con{+back,after}:      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))    case before after: tensor# The view of the padded extents that holds the tensor: on every axis, the# extent of the tensor starting after the elements added in front.def interior(before: List<&2,U32>,shape: List<&2,U32>,view: Maybe<&2,Views.View>,+axis: U32) -> Maybe<&2,Views.View>:  match before shape:    case Con{front,before} Con{extent,shape}: interior(before,shape,sliced(view,axis,front,extent,1),(axis + 1 : U32))    case before shape: viewdef all_zero(values: List<&2,U32>) -> Bool:  match values:    case []: True{}    case value <> rest: U32.is_zero(value) && all_zero(rest)# Constant padding: storage of the padded extents filled with the value, then# the tensor written through the interior view. One copy whatever the axes.def pad_constant(~Element: Data,unchanged: Bool,+value: Element,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+padded: List<&2,U32>) -> Tensors.Tensor<Element>:  match unchanged:    case True{}: tensor    case False{}:      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,part => part_shape => part_fill => part_status => part_workers =>        +part_shape = part_shape        +part_fill = part_fill        +part_workers = part_workers        Tensors.with_dense(~Element,~Tensors.Tensor<Element>,blank(~Element,Tensors.count_of(padded),padded,value,part_fill,part_workers),          target => target_shape => target_fill => target_status => target_workers =>            result_of(~Element,place_status(~Element,"pad",target_status,part_status,interior(before,part_shape,Tensors.view_of(padded),0),              target,padded,target_fill,part,part_shape,part_fill,part_workers))))def 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>:  match border padded:    case Constant{value} Some{padded}: pad_constant(~Element,all_zero(before) && all_zero(after),value,tensor,before,padded)    case border padded: pad_edge_axes(~Element,before,after,tensor,0)def padded_count(shape: Maybe<&2,List<&2,U32>>) -> Maybe<&2,U32>:  match shape:    case Some{shape}: Tensors.count_of(shape)    case None{}: None{}def pad_valid(~Element: Data,+border: Border<Element>,status: Tensors.Status,valid: Bool,tensor: Tensors.Tensor<Element>,shape: List<&2,U32>,  before: List<&2,U32>,after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor<Element>:  match status valid:    case Tensors.Usable{} True{}: pad_border(~Element,border,tensor,before,after,padded)    case Tensors.Usable{} False{}: Tensors.with_status(~Element,Tensors.Failed{Tensors.Problem{border_name(~Element,border),Tensors.InvalidArgument{shape}}},tensor)    case Tensors.Failed{problem} valid: tensor# before[i] elements are added in front of axis i and after[i] behind it. Both# lists have one entry per axis and the padded extents must be allocatable.def pad(~Element: Data,+border: Border<Element>,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+after: List<&2,U32>) -> Tensors.Tensor<Element>:  Tensors.with_metadata(~Element,~Tensors.Tensor<Element>,tensor,tensor => shape => fill => status => workers =>    +shape = shape    +padded = Shapes.padded(shape,before,after)    pad_valid(~Element,border,status,Views.has_count(padded_count(padded)),tensor,shape,before,after,padded))