def copy_result_mapping

bendtorch-core@0.1.0.0 · src/tensor/transform_PROOF.bend:188

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model

188def copy_result_mapping(-input: Shape.Shape, -axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
189  +depth: Nat, +length: Nat, array: Array<F32>, +index: U32, source: Core.Tensor<input>,
190  result: Result<&2, &2, Core.Error, F32>, balanced: Model.Balanced(F32, depth, array),
191  -valid: {Storage.bounds(depth, length) == True{} : Bool},
192  -exact: {length == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
193  allowed: {Nat.is_lt(U32.to_nat(index), length) == True{} : Bool}) ->
194  {Transforms.copy_observe(input, axes, plan,
195    Transforms.copy_read(input, axes, plan, index, source => target => (source, Done{target}),
196      Core.TTensor{Storage.Owned{depth, length, array, balanced, valid}, exact}, (source, result)), index) ==
197    Core.result_value(result) : Maybe<&2, F32>}:
198  match result:
199    case Fail{error}:
200      {==}
201    case Done{value}:
202      copy_checked_mapping(input, axes, plan, depth, length, array, index, source, value, balanced, valid, exact,
203        Nat.is_lt(U32.to_nat(index), length), {==}, {==}, allowed)