def cell_value

bendtorch-core@0.1.0.0 · src/autograd/reduce_parallel_proof.bend:778

import Base
import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
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

778# A cell of the parallel build holds `ReduceLaws.fiber`: the words its walk
779# saw add up, from +0.0, to the masked fold over the contribution.
780def cell_value(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
781  +exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>},
782  +values: List<&2, F32>, +index: U32, +d: Nat, -A: Array<F32>, +words: List<&2, F32>,
783  -seen: Obs(d, A, visited_of(rows, Shape.extents(Shape.output_shape(plan)), U32.to_nat(index)), words),
784  +inside: {Nat.is_lt(U32.to_nat(index), Shape.numel(Shape.extents(Shape.left_shape(plan)))) == True{} : Bool},
785  -reads: Reads(d, A, values, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))))) ->
786  Parallel.Fact<{Some{fsum(words, 0.0)} == ReduceLaws.fiber(plan, values, index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : Maybe<&2, F32>}>:
787  +t = U32.to_nat(index)
788  +output = Shape.extents(Shape.output_shape(plan))
789  +count = Shape.numel(output)
790  +walk = plan_walk(plan, rows, exact, t, inside)
791  +spot = masked_below(plan, rows, exact, t, count, 0n, Conversion.le_reflexive(count))
792  Parallel.Fact{
793    %Equal.sym(Maybe<&2, F32>, ReduceLaws.fiber(plan, values, index, count), gather(values, Fiber.masked(count, rows, t, 0n, 0n), Some{0.0}),
794      fiber_masked(plan, rows, exact, values, index, count, Conversion.le_reflexive(count))) :
795      {Some{fsum(words, 0.0)} == _ : Maybe<&2, F32>}
796    %Equal.sym(List<&2, Nat>, Fiber.masked(count, rows, t, 0n, 0n), visited_of(rows, output, t), walk) :
797      {Some{fsum(words, 0.0)} == gather(values, _, Some{0.0}) : Maybe<&2, F32>}
798    {\{Parallel.Fact: p => p} : Parallel.Fact<{Some{fsum(words, 0.0)} == gather(values, visited_of(rows, output, t), Some{0.0}) : Maybe<&2, F32>}> ->
799      {Some{fsum(words, 0.0)} == gather(values, visited_of(rows, output, t), Some{0.0}) : Maybe<&2, F32>}}(
800      obs_gather(d, A, values, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))), visited_of(rows, output, t), words, 0.0, seen,
801        %walk :
802          {all_below(_, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}
803        spot,
804        reads))}