def fast_fiber

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

import Base
import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

555# The fast fiber returns what the fiber returns.
556def fast_fiber(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, fuel: Nat, single: Bool, -ok: SingleOk(single, rows, extents),
557  +target: Nat, +base: Nat, +digit: Nat, +accumulator: F32, -array: Array<F32>) ->
558  Generic.Fact<{Reduce.fiber_fast(rows, extents, fuel, single, target, base, digit, accumulator, array) ==
559    Reduce.fiber(rows, extents, fuel, target, base, digit, accumulator, array) : F32 & Array<F32>}>:
560  match rows extents fuel single:
561    case Nil{} _ _ single:
562      Generic.Fact{{==}}
563    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others fuel True{}:
564      Generic.Fact{{\{Generic.Fact: p => p} : Generic.Fact<({Reduce.strided(fuel, right, base, digit, Reduce.single_offset(rest, others, Nat.mod(target, left), 0n), accumulator, array) ==
565        Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, digit, accumulator, array) : F32 & Array<F32>})> ->
566        {Reduce.strided(fuel, right, base, digit, Reduce.single_offset(rest, others, Nat.mod(target, left), 0n), accumulator, array) ==
567          Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, digit, accumulator, array) : F32 & Array<F32>}}(
568        strided_fiber(stride, left, right, rest, extent, others, ok, fuel, target, base, digit, accumulator, array))}
569    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n False{}:
570      Generic.Fact{{==}}
571    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining False{}:
572      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
573      +extents : List<&2, Nat> = extent <> others
574      +inner = Nat.mod(target, left)
575      +offset = Nat.add(base, Nat.mul(right, digit))
576      Generic.Fact{Equal.trans(F32 & Array<F32>,
577        Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k),
578          Reduce.fiber_fast(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)),
579        Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k),
580          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)),
581        Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k),
582          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)),
583        Equal.cong(F32 & Array<F32>, F32 & Array<F32>,
584          p => Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k), p),
585          Reduce.fiber_fast(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array),
586          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array),
587          {\{Generic.Fact: p => p} : Generic.Fact<({Reduce.fiber_fast(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) ==
588            Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) : F32 & Array<F32>})> ->
589            {Reduce.fiber_fast(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) ==
590              Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) : F32 & Array<F32>}}(
591            fast_fiber(rest, others, Reduce.first_span(rest, others), Reduce.single_below(rest, others),
592              single_ok(rest, others, Reduce.single_below(rest, others), {==}), inner, offset, Reduce.first_digit(rest, inner), accumulator, array))),
593        {\{Generic.Fact: p => p} : Generic.Fact<({Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k),
594          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)) ==
595          Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k),
596            Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)) : F32 & Array<F32>})> ->
597          {Reduce.walk_next(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k),
598            Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)) ==
599            Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k),
600              Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)) : F32 & Array<F32>}}(
601          next_same(v => k => Reduce.fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, v, k),
602            v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k),
603            value => kept => fast_fiber(rows, extents, remaining, False{}, Unit{}, target, base, 1n+digit, value, kept),
604            Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array))))}
605    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _ single:
606      Generic.Fact{{==}}