def run_lift

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

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/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/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/index_map.bend as IndexMap
import bendtorch-core@0.1.0.0/src/tensor/contracts/build_facts.bend as Facts
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
import bendtorch-core@0.1.0.0/src/tensor/permute_build.bend as Build
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api

210def run_lift(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +fill: F32, +ops: List<&2, StorageLaws.Op>,
211  -Q: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
212  h: Final(input, axes, plan, fill, ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), pair => Q(Transforms.copy_pack(input, axes, plan, pair))),
213  -source_fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
214  Q(Transforms.copy_allocated(input, axes, plan, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, Core.new(Layout.output(input, axes, plan), 0.0))):
215  Replayed.new_lift(Layout.output(input, axes, plan), 0.0, result => Q(Transforms.copy_allocated(input, axes, plan, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, result)),
216    target_fits => Owned.created_lift(Shape.count(Layout.output(input, axes, plan)), 0.0, Conversion.shape_count_within(Layout.output(input, axes, plan)),
217      s => (Kept(input, axes, plan, Q, Replayed.storage(input, fill, ops), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), s)) -> @-fo: {Lifecycle.logical_length(s) == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat} ->
218        Q(Transforms.copy_pack(input, axes, plan, Transforms.built(input, axes, plan, Core.TTensor{Replayed.storage(input, fill, ops), source_fits}, Core.TTensor{s, fo})))),
219      zero => hyp => fo => empty_lift(input, axes, plan, fill, ops, zero, Q, fo)(source_fits, hyp),
220      d => fits => f => hyp => fo =>
221        +positive = Arith.and_left(Nat.is_lt(0n, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))), Nat.is_le(U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Model.capacity(d)),
222          Arith.and_right(Nat.is_le(d, 20n), Bool.and(Nat.is_lt(0n, U32.to_nat(Shape.count(Layout.output(input, axes, plan)))), Nat.is_le(U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Model.capacity(d))), fits))
223        Facts.input_lift(input, fill, ops,
224          %count_same(input, axes, plan) : {Nat.is_lt(0n, _) == True{} : Bool}
225          positive,
226          st => (@-sf: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(input)) : Nat} -> Kept(input, axes, plan, Q, st, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f})) ->
227            Q(Transforms.copy_pack(input, axes, plan, Transforms.built_owned(input, axes, plan, d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), f, fo, Core.TTensor{st, sf})))),
228          ds => fits_s => b => f2 => sf => hyp2 =>
229            cloned_lift(input, axes, plan, fill, ops, d, f, fo, ds, b, f2, Array.clone(F32, Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), {==}, sf, Q,
230              kf => Facts.storage_shift(Layout.output(input, axes, plan), StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
231                Storage.Owned{d, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Build.build(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), Build.balanced(d, Build.Request{input, axes}, U32.to_nat(Shape.count(Layout.output(input, axes, plan))), Api.arr(Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill)), f},
232                made(input, axes, plan, fill, ops, d, fits, f, ds, fits_s),
233                st => ko => Packed(input, axes, plan, Q, StorageLaws.step(Storage.Owned{ds, U32.to_nat(Shape.count(input)), Api.arr(Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{}), ds, fill), b, f2}, StorageLaws.First{}), kf, st, ko),
234                ko => hyp2(kf, ko))(fo)))(source_fits, hyp))(
235      kf => ko => Facts.storage_shift(input, kept_source(input, fill, ops), StorageLaws.step(Replayed.storage(input, fill, ops), StorageLaws.First{}),
236        Replayed.replay_snoc(ops, StorageLaws.First{}, input, fill),
237        st => k => Packed(input, axes, plan, Q, st, k, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops), StorageLaws.created(Shape.count(Layout.output(input, axes, plan)), 0.0)), ko),
238        k => h(k, ko))(kf), target_fits))