def obs_gather

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

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws

753def obs_gather(+d: Nat, -A: Array<F32>, +values: List<&2, F32>, +bound: Nat, +positions: List<&2, Nat>, +words: List<&2, F32>,
754  +accumulator: F32, -seen: Obs(d, A, positions, words), +below: {all_below(positions, bound) == True{} : Bool},
755  -reads: Reads(d, A, values, bound)) ->
756  Parallel.Fact<{Some{fsum(words, accumulator)} == gather(values, positions, Some{accumulator}) : Maybe<&2, F32>}>:
757  match positions words:
758    case Nil{} Nil{}:
759      Parallel.Fact{{==}}
760    case +p <> +ps +w <> +ws:
761      +inside = Arith.and_left(Nat.is_lt(p, bound), all_below(ps, bound), below)
762      +others = Arith.and_right(Nat.is_lt(p, bound), all_below(ps, bound), below)
763      Parallel.Fact{{\{ObsStep: observed => more =>
764        %Equal.sym(Maybe<&2, F32>, ReduceLaws.word_at(values, p), Some{w},
765          Equal.trans(Maybe<&2, F32>, ReduceLaws.word_at(values, p), Model.observe(F32, d, A, Refinement.route(d, U32.from_nat(p))), Some{w},
766            Equal.sym(Maybe<&2, F32>, Model.observe(F32, d, A, Refinement.route(d, U32.from_nat(p))), ReduceLaws.word_at(values, p), reads(p, inside)),
767            observed)) :
768          {Some{fsum(ws, F32.add(accumulator, w))} == gather(values, ps, ReduceLaws.add_word(Some{accumulator}, _)) : Maybe<&2, F32>}
769        {\{Parallel.Fact: p => p} : Parallel.Fact<{Some{fsum(ws, F32.add(accumulator, w))} == gather(values, ps, Some{F32.add(accumulator, w)}) : Maybe<&2, F32>}> ->
770          {Some{fsum(ws, F32.add(accumulator, w))} == gather(values, ps, Some{F32.add(accumulator, w)}) : Maybe<&2, F32>}}(
771          obs_gather(d, A, values, bound, ps, ws, F32.add(accumulator, w), more, others, reads))} :
772        ObsStep<d, A, p, w, ps, ws> -> {Some{fsum(ws, F32.add(accumulator, w))} == gather(values, p <> ps, Some{accumulator}) : Maybe<&2, F32>}}(seen)}
773    case Nil{} +w <> +ws:
774      Parallel.Fact{Empty.absurd({Some{fsum(w <> ws, accumulator)} == gather(values, Nil{}, Some{accumulator}) : Maybe<&2, F32>}, seen)}
775    case +p <> +ps Nil{}:
776      Parallel.Fact{Empty.absurd({Some{accumulator} == gather(values, p <> ps, Some{accumulator}) : Maybe<&2, F32>}, seen)}