def parallel_allocated

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

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/elementwise.bend as Elementwise

373def parallel_allocated(+left: Shape.Shape, +right: Shape.Shape, +output: Shape.Shape, +inner: Nat, +columns: Nat,
374  a: Core.Tensor<left>, b: Core.Tensor<right>, result: Result<&2, &1, Core.Error, Core.Tensor<output>>) ->
375  Result<&2, &1, Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>:
376  match result:
377    case Fail{error}:
378      Fail{TensorError{error}}
379    case Done{Core.TTensor{storage, exact}}:
380      match storage:
381        case Storage.Empty{}:
382          Done{(Core.pack(left, a), Core.pack(right, b), Core.pack(output, Core.TTensor{Storage.Empty{}, exact}))}
383        case Storage.Owned{+depth, +length, array, balanced, valid}:
384          parallel_output(output, inner, columns, depth, length, valid, exact,
385            Elementwise.take(Core.pack(left, a)), Elementwise.take(Core.pack(right, b)))