def parallel_output

bendtorch-core@0.1.0.0 · src/tensor/matmul.bend:359

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

359# Both inputs come back to the caller as the kept halves of their clones,
360# and the build reads the other halves.
361def parallel_output(+output: Shape.Shape, +inner: Nat, +columns: Nat, +depth: Nat, +length: Nat,
362  -valid: {Storage.bounds(depth, length) == True{} : Bool},
363  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
364  left: Core.Dynamic & Array<F32>, right: Core.Dynamic & Array<F32>) ->
365  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
366  (kept_left, la) = left
367  (kept_right, ra) = right
368  Done{(kept_left, kept_right, Core.pack(output, Core.TTensor{Storage.Owned{depth, length,
369    build(depth, product_splits(depth, inner), inner, columns, length, 0n, la, ra),
370    build_balanced(depth, product_splits(depth, inner), inner, columns, length, 0n, la, ra),
371    valid}, exact}))}