def selected_leaf

stelliferous@0.0.2.0 · tensor_select.bend:42

import Base
import stelliferous@0.0.2.0/tensor_lanes.bend as Lanes
import stelliferous@0.0.2.0/tensor_band.bend as Bands
import stelliferous@0.0.2.0/storage_buffer.bend as Storage
import stelliferous@0.0.2.0/traversal_partition.bend as Partition

42def selected_leaf(geometry: Bands.Geometry,first: U32,rows: U32,
43  input: Bands.Paired<Storage.Buffer<F32>,Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>>,resource: Unit) -> Bands.Bands(~F32):
44  match input:
45    case Bands.Both{yes,Bands.Both{mask,no}}:
46      +rows = rows
47      Partition.Leaf{Lanes.select(U32.to_nat(Bands.local_count(geometry,rows)),mask,yes,no),first,rows,Unit{}}
48    case Bands.Both{yes,Bands.Only{mask,own_first,own_rows}}: Partition.Leaf{yes,first,rows,Unit{}}
49    case Bands.Only{yes,own_first,own_rows}: Partition.Leaf{yes,own_first,own_rows,Unit{}}