def fiber_masked

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

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/tensor/projection_LAWS.bend as Projection
import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast_proof.bend as BroadcastProof
import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
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

140# The masked fold of `ReduceLaws.fiber` adds, in order, the words at the
141# masked positions.
142def fiber_masked(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
143  +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>},
144  +values: List<&2, F32>, +index: U32, +count: Nat,
145  +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) ->
146  {ReduceLaws.fiber(plan, values, index, count) == gather(values, Fiber.masked(count, rows, U32.to_nat(index), 0n, 0n), Some{0.0}) : Maybe<&2, F32>}:
147  match count:
148    case 0n:
149      {==}
150    case 1n+k:
151      +t = U32.to_nat(index)
152      +map : Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))> =
153        Shape.Indexed{rows, exact}
154      +limit = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))
155      +below = Conversion.le_successor_lt(k, limit, room)
156      +lm = Shape.walk_left(rows, k)
157      +rm = Shape.walk_right(rows, k)
158      +before = Fiber.masked(k, rows, t, 0n, 0n)
159      +last = Fiber.kept(Nat.is_eq(lm, t), rm, Nil{})
160      +lm_below = BroadcastProof.left_walk_bound(plan, map, k, below)
161      +rm_below = BroadcastProof.right_walk_bound(plan, map, k, below)
162      %Equal.sym(Maybe<&1, U32 & U32>, Projection.offset_spec(plan, k), Some{(U32.from_nat(lm), U32.from_nat(rm))},
163        offsets_walk(plan, map, k, BroadcastProof.output_inside(plan, k, below))) :
164        {ReduceLaws.contribute(values, index, ReduceLaws.fiber(plan, values, index, k), _) ==
165          gather(values, Fiber.masked(1n+k, rows, t, 0n, 0n), Some{0.0}) : Maybe<&2, F32>}
166      %Equal.sym(Nat, U32.to_nat(U32.from_nat(rm)), rm, PointwiseProof.word_exact(rm, Shape.right_shape(plan), rm_below)) :
167        {ReduceLaws.matched(ReduceLaws.fiber(plan, values, index, k), ReduceLaws.word_at(values, _), Nat.is_eq(U32.to_nat(U32.from_nat(lm)), t)) ==
168          gather(values, Fiber.masked(1n+k, rows, t, 0n, 0n), Some{0.0}) : Maybe<&2, F32>}
169      %Equal.sym(Nat, U32.to_nat(U32.from_nat(lm)), lm, PointwiseProof.word_exact(lm, Shape.left_shape(plan), lm_below)) :
170        {ReduceLaws.matched(ReduceLaws.fiber(plan, values, index, k), ReduceLaws.word_at(values, rm), Nat.is_eq(_, t)) ==
171          gather(values, Fiber.masked(1n+k, rows, t, 0n, 0n), Some{0.0}) : Maybe<&2, F32>}
172      %Equal.sym(Maybe<&2, F32>, ReduceLaws.fiber(plan, values, index, k), gather(values, before, Some{0.0}),
173        fiber_masked(plan, rows, exact, values, index, k, Conversion.le_predecessor(k, limit, room))) :
174        {ReduceLaws.matched(_, ReduceLaws.word_at(values, rm), Nat.is_eq(lm, t)) ==
175          gather(values, Fiber.masked(1n+k, rows, t, 0n, 0n), Some{0.0}) : Maybe<&2, F32>}
176      %succ_add(k) :
177        {ReduceLaws.matched(gather(values, before, Some{0.0}), ReduceLaws.word_at(values, rm), Nat.is_eq(lm, t)) ==
178          gather(values, Fiber.masked(_, rows, t, 0n, 0n), Some{0.0}) : Maybe<&2, F32>}
179      %Equal.sym(List<&2, Nat>, Fiber.masked(Nat.add(k, 1n), rows, t, 0n, 0n), List.append(&2, Nat, before, Fiber.masked(1n, rows, t, 0n, Nat.add(0n, k))),
180        Fiber.masked_split(k, 1n, rows, t, 0n, 0n)) :
181        {ReduceLaws.matched(gather(values, before, Some{0.0}), ReduceLaws.word_at(values, rm), Nat.is_eq(lm, t)) ==
182          gather(values, _, Some{0.0}) : Maybe<&2, F32>}
183      %Equal.sym(Maybe<&2, F32>, gather(values, List.append(&2, Nat, before, last), Some{0.0}), gather(values, last, gather(values, before, Some{0.0})),
184        gather_append(values, before, last, Some{0.0})) :
185        {ReduceLaws.matched(gather(values, before, Some{0.0}), ReduceLaws.word_at(values, rm), Nat.is_eq(lm, t)) == _ : Maybe<&2, F32>}
186      Equal.sym(Maybe<&2, F32>, gather(values, last, gather(values, before, Some{0.0})),
187        ReduceLaws.matched(gather(values, before, Some{0.0}), ReduceLaws.word_at(values, rm), Nat.is_eq(lm, t)),
188        gather_kept(values, Nat.is_eq(lm, t), rm, gather(values, before, Some{0.0})))