def masked_below
bendtorch-core@0.1.0.0 · src/autograd/reduce_parallel_proof.bend:208
import bendtorch-core@0.1.0.0/src/autograd/reduce_parallel_proof.bend as ReduceParallelProof
import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast_proof.bend as BroadcastProof
import bendtorch-core@0.1.0.0/src/autograd/reduce_fiber_proof.bend as Fiber
208def masked_below(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
209 +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>},
210 +target: Nat, +n: Nat, +cursor: Nat,
211 +room: {Nat.is_le(Nat.add(cursor, n), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
212 {all_below(Fiber.masked(n, rows, target, 0n, cursor), Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}:
213 match n:
214 case 0n:
215 {==}
216 case 1n+m:
217 +map : Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))> =
218 Shape.Indexed{rows, exact}
219 kept_below(Nat.is_eq(Shape.walk_left(rows, cursor), target), Shape.walk_right(rows, cursor), Fiber.masked(m, rows, target, 0n, 1n+cursor),
220 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))),
221 BroadcastProof.right_walk_bound(plan, map, cursor, Fiber.lt_of_room(cursor, m, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), room)),
222 masked_below(plan, rows, exact, target, m, 1n+cursor, Fiber.le_step(cursor, m, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), room)))