def parallel_retains

bendtorch-core@0.1.0.0 · src/autograd/reduce_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/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/autograd/reduce.bend as Reduce
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as Laws

188def parallel_retains(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
189  result: Result<&2, &1, Core.Error, Core.Tensor<target>>) ->
190  Laws.Retained<Core.Packed{shape, tensor}, Reduce.parallel_allocated(Core.Packed{shape, tensor}, target, plan, result)>:
191  match result:
192    case Fail{error}:
193      Laws.Retained{property => {==}}
194    case Done{Core.TTensor{storage, exact}}:
195      match storage:
196        case Storage.Empty{}:
197          Laws.Retained{property => {==}}
198        case Storage.Owned{+depth, +length, array, balanced, valid}:
199          indexed_retains(shape, tensor, target, plan, Core.TTensor{Storage.Owned{depth, length, array, balanced, valid}, exact}, depth, length,
200            valid, exact, Shape.index_map(plan))