def reduce_fold_fiber

bendtorch-core@0.1.0.0 · src/autograd/reduce_fold_PROOF.bend:237

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/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
import bendtorch-core@0.1.0.0/src/tensor/contracts/fold_proof.bend as FoldProof
import bendtorch-core@0.1.0.0/src/tensor/contracts/sum.bend as Sum
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_fold.bend as ReduceFold
import bendtorch-core@0.1.0.0/src/autograd/reduce_fold_LAWS.bend as Laws

237def reduce_fold_fiber(+plan: Shape.BroadcastPlan,
238  same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>},
239  +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>,
240  -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) ==
241    Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}) ->
242  Laws.FiberContract<Shape.right_shape(plan),
243    ReduceLaws.fiber(plan,
244      Contracts.words(Shape.right_shape(plan), fill, ops, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
245      index, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))),
246    other => Replayed.value(Shape.right_shape(plan), fill, ops, other),
247    ReduceFold.cell(plan, index,
248      Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>:
249  +right = Shape.right_shape(plan)
250  +count = Shape.elements(right)
251  +values = Contracts.words(right, fill, ops, Shape.logical_numel(Shape.dimensions(right)))
252  +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index}
253  +total = Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, count)
254  +room = contribution_room(plan, same)
255  fiber_contract(plan, index, fill, ops, fits,
256    ReduceLaws.fiber(plan, values, index, Shape.logical_numel(Shape.dimensions(right))), total, count,
257    FoldProof.contract(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, count,
258      source_inside(plan, index, count, room), fits),
259    %Shape.elements_exact(right) : {ReduceLaws.fiber(plan, values, index, _) == Some{total} : Maybe<&2, F32>}
260    fiber_total(plan, index, fill, ops, count, room))