def cell_fiber

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

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/contracts/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
import bendtorch-core@0.1.0.0/src/autograd/reduce_replayed_proof.bend as Replay

193# A target cell of the replayed reduce holds `ReduceLaws.fiber` over the
194# contribution's words.
195def cell_fiber(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan,
196  +planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
197  +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>},
198  +fill: F32, +ops: List<&2, StorageLaws.Op>, +index: U32,
199  +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool}) ->
200  {Some{Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index))} ==
201    ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index,
202      Shape.logical_numel(Shape.dimensions(shape))) : Maybe<&2, F32>}:
203  +left = plan_left(target, shape, plan, planned)
204  %contribution_count(shape, plan, same) :
205    {Some{Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index))} ==
206      ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index, _) : Maybe<&2, F32>}
207  Replay.word_fiber(plan, Shape.index_map(plan), shape, fill, ops, index,
208    %Equal.sym(Shape.Shape, Shape.left_shape(plan), target, left) : {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(_))) == True{} : Bool}
209    inside,
210    Equal.cong(Shape.Shape, List<&2, U32>, Shape.dimensions, Shape.right_shape(plan), shape, plan_right(target, shape, plan, planned)))