def rfold_gather

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

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/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
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_parallel_proof.bend as ReduceParallel

1077def rfold_gather(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +positions: List<&2, Nat>, +accumulator: F32,
1078  +below: {ReduceParallel.all_below(positions, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
1079  {Some{rfold(shape, fill, ops, positions, accumulator)} ==
1080    ReduceParallel.gather(Laws.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), positions, Some{accumulator}) : Maybe<&2, F32>}:
1081  match positions:
1082    case Nil{}:
1083      {==}
1084    case +p <> +ps:
1085      +count = Shape.logical_numel(Shape.dimensions(shape))
1086      +here = Arith.and_left(Nat.is_lt(p, count), ReduceParallel.all_below(ps, count), below)
1087      +later = Arith.and_right(Nat.is_lt(p, count), ReduceParallel.all_below(ps, count), below)
1088      +word = Replayed.value(shape, fill, ops, U32.from_nat(p))
1089      %Equal.sym(Maybe<&2, F32>, ReduceLaws.word_at(Laws.words(shape, fill, ops, count), p), Some{word}, words_at(shape, fill, ops, count, p, here)) :
1090        {Some{rfold(shape, fill, ops, ps, F32.add(accumulator, word))} ==
1091          ReduceParallel.gather(Laws.words(shape, fill, ops, count), ps, ReduceLaws.add_word(Some{accumulator}, _)) : Maybe<&2, F32>}
1092      rfold_gather(shape, fill, ops, ps, F32.add(accumulator, word), later)