def run_lift

bendtorch-core@0.1.0.0 · src/autograd/reduce_to_fold_PROOF.bend:151

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/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
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/autograd/reduce.bend as Reduce
import bendtorch-core@0.1.0.0/src/tensor/read_contract.bend as Reads
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
import bendtorch-core@0.1.0.0/src/autograd/reduce_replayed_proof.bend as Replay

151# The parallel reduce over a replayed contribution, in the transport form of
152# `Replay.Final`.
153def run_lift(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan,
154  +planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
155  +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>},
156  +fill: F32, +ops: List<&2, StorageLaws.Op>,
157  -P: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic> -> Type,
158  h: Replay.Final(plan, target, shape, fill, ops, Shape.index_map(plan), StorageLaws.created(Shape.count(target), 0.0), P)) ->
159  @-fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
160    P(Reduce.reduce_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)):
161  +output = Shape.output_shape(plan)
162  +left = plan_left(target, shape, plan, planned)
163  +right = plan_right(target, shape, plan, planned)
164  +count_left : {U32.to_nat(Shape.count(target)) == Shape.numel(Shape.extents(Shape.left_shape(plan))) : Nat} =
165    %Equal.sym(Shape.Shape, Shape.left_shape(plan), target, left) :
166      {U32.to_nat(Shape.count(target)) == Shape.logical_numel(Shape.dimensions(_)) : Nat}
167    Reads.count_exact(target)
168  +facts : Replay.Facts<plan, shape, U32.to_nat(Shape.count(target))> =
169    Replay.Facts{count_left, Equal.cong(Shape.Shape, List<&2, U32>, Shape.dimensions, Shape.right_shape(plan), shape, right)}
170  fits =>
171    %Equal.sym(Result<&2, &2, Shape.Error, Shape.BroadcastPlan>, Shape.broadcast_checked(target, shape), Done{plan}, planned) :
172      P(Reduce.planned(shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}, target, _))
173    %Equal.sym(Bool, Layout.words_same(Shape.dimensions(output), Shape.dimensions(shape)), True{}, fitted_same(shape, plan, same)) :
174      P(Reduce.fitted(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan, _))
175    Replayed.new_lift(target, 0.0,
176      result => P(Reduce.parallel_allocated(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan, result)),
177      fo => Owned.created_lift(Shape.count(target), 0.0, Conversion.shape_count_within(target),
178        st => (Replay.Final(plan, target, shape, fill, ops, Shape.index_map(plan), st, P) ->
179          @-fo: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(target)) : Nat} ->
180          P(Reduce.parallel_allocated(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan, Done{Core.TTensor{st, fo}}))),
181        zero => final => fo =>
182          final(ops)(other => {==})({==})(fits)(Nil{})(
183            index => inside => Empty.absurd(
184              {Replayed.value(target, 0.0, Nil{}, index) == Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index)) : F32},
185              Arith.lt_zero(U32.to_nat(index),
186                %Equal.trans(Nat, Shape.logical_numel(Shape.dimensions(target)), U32.to_nat(Shape.count(target)), 0n,
187                  Equal.sym(Nat, U32.to_nat(Shape.count(target)), Shape.logical_numel(Shape.dimensions(target)), Reads.count_exact(target)), zero) :
188                  {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool}
189                inside)))(fo),
190        d => fits_d => f => final => fo =>
191          Replay.indexed_lift(plan, target, shape, fill, ops, fits, facts, d, fits_d, f, fo, Shape.index_map(plan), P, final))(h)(fo))