def word_fiber
bendtorch-core@0.1.0.0 · src/autograd/reduce_replayed_proof.bend:1094
import bendtorch-core@0.1.0.0/src/autograd/reduce_replayed_proof.bend as ReduceReplayedProof
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/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
import bendtorch-core@0.1.0.0/src/autograd/reduce_fiber_proof.bend as Fiber
import bendtorch-core@0.1.0.0/src/autograd/reduce_parallel_proof.bend as ReduceParallel
1094def word_fiber(+plan: Shape.BroadcastPlan, +map: PlanMap(plan), +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +index: U32,
1095 +inside: {Nat.is_lt(U32.to_nat(index), Shape.numel(Shape.extents(Shape.left_shape(plan)))) == True{} : Bool},
1096 +right: {Shape.dimensions(Shape.right_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}) ->
1097 {Some{cell_word(plan, map, shape, fill, ops, U32.to_nat(index))} ==
1098 ReduceLaws.fiber(plan, Laws.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index,
1099 Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : Maybe<&2, F32>}:
1100 match map:
1101 case Shape.Unindexed{zero}:
1102 %zero : {Some{0.0} == ReduceLaws.fiber(plan, Laws.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index, _) : Maybe<&2, F32>}
1103 {==}
1104 case Shape.Indexed{+rows, +rows_exact}:
1105 +t = U32.to_nat(index)
1106 +output = Shape.extents(Shape.output_shape(plan))
1107 +count = Shape.numel(output)
1108 +values = Laws.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape)))
1109 +walk = ReduceParallel.plan_walk(plan, rows, rows_exact, t, inside)
1110 %Equal.sym(Maybe<&2, F32>, ReduceLaws.fiber(plan, values, index, count), ReduceParallel.gather(values, Fiber.masked(count, rows, t, 0n, 0n), Some{0.0}),
1111 ReduceParallel.fiber_masked(plan, rows, rows_exact, values, index, count, Conversion.le_reflexive(count))) :
1112 {Some{rfold(shape, fill, ops, ReduceParallel.visited_of(rows, output, t), 0.0)} == _ : Maybe<&2, F32>}
1113 %Equal.sym(List<&2, Nat>, Fiber.masked(count, rows, t, 0n, 0n), ReduceParallel.visited_of(rows, output, t), walk) :
1114 {Some{rfold(shape, fill, ops, ReduceParallel.visited_of(rows, output, t), 0.0)} == ReduceParallel.gather(values, _, Some{0.0}) : Maybe<&2, F32>}
1115 rfold_gather(shape, fill, ops, ReduceParallel.visited_of(rows, output, t), 0.0,
1116 walk_below(plan, rows, rows_exact, shape, Shape.numel(Shape.extents(Shape.left_shape(plan))), Facts{{==}, right}, t, inside))