def built_balanced

bendtorch-core@0.1.0.0 · src/tensor/gather.bend:369

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice

369def built_balanced(kind: Pass, +depth: Nat, +plan: Plan, +count: Nat, a: Array<F32>, index: Array<U32>) ->
370  Model.Balanced(F32, depth, built_array(kind, depth, plan, count, a, index)):
371  match kind:
372    case Forward{}:
373      balanced(depth, Slice.splits(depth), plan, count, 0n, a, index)
374    case Backward{}:
375      scatter_balanced(depth, Slice.splits(depth), plan, count, 0n, a, index)