def walk_below

bendtorch-core@0.1.0.0 · src/autograd/reduce_replayed_proof.bend:714

import Base
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_parallel_proof.bend as ReduceParallel

714# Every offset the walk of a cell below `count` reads lies inside the
715# contribution.
716def walk_below(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
717  +rows_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>},
718  +shape: Shape.Shape, +count: Nat, +facts: Facts<plan, shape, count>, +base: Nat, +inside: {Nat.is_lt(base, count) == True{} : Bool}) ->
719  {ReduceParallel.all_below(ReduceParallel.visited_of(rows, Shape.extents(Shape.output_shape(plan)), base), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}:
720  Facts{+count_left, +right} = facts
721  +output = Shape.extents(Shape.output_shape(plan))
722  +below_left : {Nat.is_lt(base, Shape.numel(Shape.extents(Shape.left_shape(plan)))) == True{} : Bool} =
723    %count_left : {Nat.is_lt(base, _) == True{} : Bool}
724    inside
725  %right : {ReduceParallel.all_below(ReduceParallel.visited_of(rows, output, base), Shape.logical_numel(_)) == True{} : Bool}
726  %ReduceParallel.plan_walk(plan, rows, rows_exact, base, below_left) :
727    {ReduceParallel.all_below(_, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}
728  ReduceParallel.masked_below(plan, rows, rows_exact, base, Shape.numel(output), 0n, Conversion.le_reflexive(Shape.numel(output)))