def max_rows

bendtorch-core@0.1.0.0 · src/autograd/loss.bend:138

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
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix

138def max_rows(+target: Shape.Shape, +columns: Nat, +depth: Nat, +length: Nat,
139  -valid: {Storage.bounds(depth, length) == True{} : Bool},
140  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
141  taken: Core.Dynamic & Array<F32>) -> Core.Dynamic & Result<&2, &1, Core.Error, Core.Dynamic>:
142  (kept, array) = taken
143  +splits = Matrix.splits(depth, columns)
144  (kept, Done{Core.Packed{target, Core.TTensor{Storage.Owned{depth, length,
145    max_build(depth, splits, columns, length, 0n, array),
146    max_balanced(depth, splits, columns, length, 0n, array),
147    valid}, exact}}})