def parallel_rows

bendtorch-core@0.1.0.0 · src/autograd/reduce.bend:251

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core

251def parallel_rows(+plan: Shape.BroadcastPlan, +target: Shape.Shape, +depth: Nat, +length: Nat,
252  -valid: {Storage.bounds(depth, length) == True{} : Bool},
253  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
254  taken: Core.Dynamic & Array<F32>, rows: List<&2, Shape.Row>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
255  (kept, array) = taken
256  +extents = Shape.extents(Shape.output_shape(plan))
257  +splits = splits(plan, depth, length)
258  (kept, Done{Core.Packed{target, Core.TTensor{Storage.Owned{depth, length,
259    build(depth, splits, rows, extents, length, 0n, array),
260    build_balanced(depth, splits, rows, extents, length, 0n, array),
261    valid}, exact}}})