def run_lift

bendtorch-core@0.1.0.0 · src/autograd/broadcast_pointwise_PROOF.bend:109

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/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/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
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_LAWS.bend as BinaryLaws
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_proofs.bend as Binary
import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
import bendtorch-core@0.1.0.0/src/autograd/reduce_to_fold_PROOF.bend as ReduceTo
import bendtorch-core@0.1.0.0/src/autograd/broadcast_pointwise_LAWS.bend as Laws

109# broadcast_to on a replayed input whose shape broadcasts to exactly the
110# target: the planner and the target check pass, the zero tensor of the target
111# is allocated, and the parallel build runs KeepLeft with the input on the
112# left and the zero tensor on the right, whose words are `cells`.
113def run_lift(+shape: Shape.Shape, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
114  +planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
115  +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
116  +fill: F32, +ops: List<&2, StorageLaws.Op>,
117  -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type,
118  h: Laws.BroadcastWith(plan, shape, fill, ops, P)) ->
119  @-fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
120    P(Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)):
121  +output = Shape.output_shape(plan)
122  +count = Shape.logical_numel(Shape.dimensions(output))
123  fits =>
124    -input = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}} : Core.Dynamic}
125    %Equal.sym(Result<&2, &2, Shape.Error, Shape.BroadcastPlan>, Shape.broadcast_checked(shape, target), Done{plan}, planned) :
126      P(Broadcast.planned(input, target, _))
127    %Equal.sym(Bool, Layout.words_same(Shape.dimensions(output), Shape.dimensions(target)), True{}, ReduceTo.fitted_same(target, plan, same)) :
128      P(Broadcast.fitted(input, target, plan, _))
129    Replayed.new_lift(target, 0.0, result => P(Broadcast.zeroed(input, target, plan, result)),
130      zeros_fits => Binary.allocated_replayed(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, planned,
131        pair => P(Broadcast.finish(pair)),
132        kept => kept_words => input_fits => zeros_kept => zeros_words => zeros_kept_fits => output_ops => values => output_fits =>
133          h(kept, kept_words, input_fits, output_ops,
134            Equal.trans(List<&2, F32>, Contracts.words(output, 0.0, output_ops, count),
135              BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, count), Laws.cells(plan, shape, fill, ops, count),
136              values, cells_keep(plan, shape, fill, ops, target, count)),
137            output_fits))(fits, zeros_fits))