def dispatch

stelliferous@0.0.2.0 · tensor_select.bend:58

import Base
import stelliferous@0.0.2.0/tensor.bend as Tensors
import stelliferous@0.0.2.0/tensor_band_schedule.bend as Schedule

58def dispatch(layout: Maybe<&2,BandPlan>,shape: List<&2,U32>,count: U32,fill: F32,workers: U32,
59  mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
60  match layout:
61    case None{}: dense(shape,count,workers,mask,yes,no)
62    case Some{BandPlan{+geometry,+layout}}:
63      +workers = workers
64      Tensors.Banded{selected_bands(geometry,layout,Schedule.parallel(geometry,workers,16.0,0.875,layout),
65        Tensors.aligned(~F32,geometry,layout,mask),Tensors.aligned(~F32,geometry,layout,yes),Tensors.aligned(~F32,geometry,layout,no)),
66        geometry,shape,fill,Tensors.Usable{},workers}