def input_take
bendtorch-core@0.1.0.0 · src/tensor/contracts/matmul_parallel_proof.bend:831
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as MatmulParallelProof
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/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/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/from_values_contract.bend as FromValues
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
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
831def input_take(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, -Q: Core.Dynamic & Array<F32> -> Type,
832 h: InputAt(shape, fill, ops, StorageLaws.created(Shape.count(shape), fill), Q)) ->
833 @-fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
834 Q(Elementwise.take(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}})):
835 Owned.created_lift(Shape.count(shape), fill, Conversion.shape_count_within(shape),
836 s => (InputAt(shape, fill, ops, s, Q) -> @-fits: {Lifecycle.logical_length(StorageLaws.replay(ops, s)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
837 Q(Elementwise.take(Core.Packed{shape, Core.TTensor{StorageLaws.replay(ops, s), fits}}))),
838 zero => k => fits =>
839 storage_shift(shape, Storage.Empty{}, StorageLaws.replay(ops, Storage.Empty{}),
840 Equal.sym(Storage.Storage, StorageLaws.replay(ops, Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(ops)),
841 st => kk => Q(Elementwise.take(Core.Packed{shape, Core.TTensor{st, kk}})),
842 kk => storage_shift(shape, StorageLaws.replay(ops, Storage.Empty{}), Storage.Empty{}, Owned.replay_empty(ops),
843 st => k2 => Q((Core.Packed{shape, Core.TTensor{st, k2}}, Array.new(F32, 0n, 0.0))),
844 k2 => k(ops)(other => {==})({==})(Nil{})(0n)(0.0)(NoCells{zero})(k2))(kk))(fits),
845 d => fits_d => f => k => fits =>
846 Owned.replay_lift(ops, d, U32.to_nat(Shape.count(shape)), fill, Nil{}, Api.hit_of(U32.to_nat(Shape.count(shape)), ops), {==},
847 ArrayCore.new_balanced(F32, d, fill), f, fits_d,
848 st => (@-fits: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
849 (@-kk: {Lifecycle.logical_length(StorageLaws.step(st, StorageLaws.First{})) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
850 Q((Core.Packed{shape, Core.TTensor{StorageLaws.step(st, StorageLaws.First{}), kk}}, Api.arr(history(shape, ops), d, fill)))) ->
851 Q(Elementwise.take(Core.Packed{shape, Core.TTensor{st, fits}}))),
852 c => fits2 => kc =>
853 take_owned_lift(shape, d, U32.to_nat(Shape.count(shape)), Api.arr(Owned.norm(U32.to_nat(Shape.count(shape)), ops, Nil{}), d, fill),
854 Array.clone(F32, Api.arr(Owned.norm(U32.to_nat(Shape.count(shape)), ops, Nil{}), d, fill)), {==}, c, f, fits2, Q, kc))(fits)(
855 storage_shift(shape,
856 StorageLaws.replay(Replayed.snoc(ops, StorageLaws.First{}), Storage.Owned{d, U32.to_nat(Shape.count(shape)), Array.new(F32, d, fill), ArrayCore.new_balanced(F32, d, fill), f}),
857 StorageLaws.step(StorageLaws.replay(ops, Storage.Owned{d, U32.to_nat(Shape.count(shape)), Array.new(F32, d, fill), ArrayCore.new_balanced(F32, d, fill), f}), StorageLaws.First{}),
858 FromValues.replay_snoc(ops, StorageLaws.First{}, Storage.Owned{d, U32.to_nat(Shape.count(shape)), Array.new(F32, d, fill), ArrayCore.new_balanced(F32, d, fill), f}),
859 st => kk => Q((Core.Packed{shape, Core.TTensor{st, kk}}, Api.arr(history(shape, ops), d, fill))),
860 k(Replayed.snoc(ops, StorageLaws.First{}))(other => half_keeps(shape, fill, ops, other))(words_first(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))))(
861 history(shape, ops))(d)(fill)(Normal{fits_d, {==}, {==}}))))(h)