def fiber_kept

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

649# A fiber of `Api.arr(h, d, f)` hands that array back.
650def fiber_kept(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat, +accumulator: F32,
651  +h: List<&2, Api.Effect>, +d: Nat, +f: F32) ->
652  {Reduce.fiber(rows, extents, fuel, target, base, digit, accumulator, Api.arr(h, d, f)) ==
653    (Pair.fst(F32, Array<F32>, Reduce.fiber(rows, extents, fuel, target, base, digit, accumulator, Api.arr(h, d, f))), Api.arr(h, d, f)) : F32 & Array<F32>}:
654  match rows extents fuel:
655    case Nil{} _ _:
656      %Equal.sym(Array<F32> & F32, Array.get(F32, Api.arr(h, d, f), U32.from_nat(base)), (Api.arr(h, d, f), word_at(h, d, f, base)), got(h, d, f, U32.from_nat(base))) :
657        {Reduce.leaf_add(accumulator, _) == (Pair.fst(F32, Array<F32>, Reduce.leaf_add(accumulator, _)), Api.arr(h, d, f)) : F32 & Array<F32>}
658      {==}
659    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
660      {==}
661    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
662      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
663      +extents : List<&2, Nat> = extent <> others
664      +inner = Nat.mod(target, left)
665      +offset = Nat.add(base, Nat.mul(right, digit))
666      +value = Pair.fst(F32, Array<F32>, Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)))
667      %Equal.sym(F32 & Array<F32>, Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)), (value, Api.arr(h, d, f)),
668        fiber_kept(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, h, d, f)) :
669        {Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k), _) ==
670          (Pair.fst(F32, Array<F32>, Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k), _)), Api.arr(h, d, f)) :
671          F32 & Array<F32>}
672      fiber_kept(rows, extents, remaining, target, base, 1n+digit, value, h, d, f)
673    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
674      {==}