def fast_arr

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

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
import bendtorch-core@0.1.0.0/src/autograd/reduce_parallel_proof.bend as ReduceParallel

676def fast_arr(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +fuel: Nat, single: Bool, +ok: ReduceParallel.SingleOk(single, rows, extents),
677  +target: Nat, +base: Nat, +digit: Nat, +accumulator: F32, +h: List<&2, Api.Effect>, +d: Nat, +f: F32) ->
678  {Reduce.fiber_fast(rows, extents, fuel, single, target, base, digit, accumulator, Api.arr(h, d, f)) ==
679    Reduce.fiber(rows, extents, fuel, target, base, digit, accumulator, Api.arr(h, d, f)) : F32 & Array<F32>}:
680  match rows extents fuel single:
681    case Nil{} _ _ single:
682      {==}
683    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others fuel True{}:
684      strided_arr(stride, left, right, rest, extent, others, ok, fuel, target, base, digit, accumulator, h, d, f)
685    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n False{}:
686      {==}
687    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining False{}:
688      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
689      +extents : List<&2, Nat> = extent <> others
690      +inner = Nat.mod(target, left)
691      +offset = Nat.add(base, Nat.mul(right, digit))
692      +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)))
693      %Equal.sym(F32 & Array<F32>,
694        Reduce.fiber_fast(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)),
695        Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)),
696        fast_arr(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others),
697          ReduceParallel.single_ok(rest, others, Reduce.single_below(rest, others), {==}), inner, offset, Reduce.first_digit(rest, inner), accumulator, h, d, f)) :
698        {Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k), _) ==
699          Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k), Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f))) : F32 & Array<F32>}
700      %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)),
701        fiber_kept(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, h, d, f)) :
702        {Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k), _) ==
703          Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k), _) : F32 & Array<F32>}
704      fast_arr(rows, extents, remaining, False{}, Unit{}, target, base, 1n+digit, value, h, d, f)
705    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _ single:
706      {==}