def planned_lift

bendtorch-core@0.1.0.0 · src/tensor/permute_values.bend:240

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed

240def planned_lift(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
241  result: Result<&2, &2, Layout.Error, Layout.Plan<input, axes>>,
242  done: {Result.is_done(&2, &2, Layout.Error, Layout.Plan<input, axes>, result) == True{} : Bool},
243  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
244  h: @+plan: Layout.Plan<input, axes> -> Final(input, axes, plan, fill, ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
245    pair => Q(Transforms.copy_pack(input, axes, plan, pair))),
246  -source_fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
247  Q(Transforms.copy_planned(input, axes, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, result)):
248  match result:
249    case Fail{error}:
250      Empty.absurd(Q(Transforms.copy_planned(input, axes, Core.TTensor{Replayed.storage(input, fill, ops), source_fits},
251        Fail{error})), Arith.clash(done))
252    case Done{+plan}:
253      run_lift(input, axes, plan, fill, ops, Q, h(plan), source_fits)