def reduce_axis_count

stelliferous@0.0.2.0 · tensor_layout.bend:449

import Base
import stelliferous@0.0.2.0/storage_buffer.bend as Storage
import stelliferous@0.0.2.0/tensor_shape.bend as Shapes
import stelliferous@0.0.2.0/tensor.bend as Tensors

449# unit is the shape with extent 1 at the axis: the layout of the output and
450# the extents of every slice.
451def reduce_axis_count(~Element: Data,~operation: Element -> Element -> Element,name: String,count: Maybe<&2,U32>,+axis: U32,input: Storage.Buffer<Element>,
452  +shape: List<&2,U32>,unit: List<&2,U32>,+initial: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
453  match count:
454    case Some{count}:
455      folded(~Element,name,shape,Shapes.without_axis(shape,U32.to_nat(axis)),fill,workers,
456        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{})))
457    case None{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},Shapes.without_axis(shape,U32.to_nat(axis)),fill,workers)