def step_arr

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

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

563def step_arr(+stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +extent: Nat, +others: List<&2, Nat>,
564  +single: {Reduce.single_rows(rest, others) == True{} : Bool}, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat, +accumulator: F32,
565  +h: List<&2, Api.Effect>, +d: Nat, +f: F32) ->
566  {Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 1n+fuel, target, base, digit, accumulator, Api.arr(h, d, f)) ==
567    Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit,
568      F32.add(accumulator, word_at(h, d, f, Reduce.strided_offset(right, base, digit, Reduce.single_offset(rest, others, Nat.mod(target, left), 0n)))),
569      Api.arr(h, d, f)) : F32 & Array<F32>}:
570  +inner = Nat.mod(target, left)
571  +lift = Reduce.single_offset(rest, others, inner, 0n)
572  +offset = Nat.add(base, Nat.mul(right, digit))
573  +at = Reduce.strided_offset(right, base, digit, lift)
574  Equal.trans(F32 & Array<F32>,
575    Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
576      Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f))),
577    Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
578      Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset))))),
579    Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
580      Reduce.leaf_add(accumulator, (Api.arr(h, d, f), word_at(h, d, f, at)))),
581    Equal.cong(F32 & Array<F32>, F32 & Array<F32>,
582      p => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k), p),
583      Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)),
584      Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))),
585      single_arr(rest, others, single, inner, offset, accumulator, h, d, f)),
586    Equal.trans(F32 & Array<F32>,
587      Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
588        Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset))))),
589      Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
590        Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(at)))),
591      Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
592        Reduce.leaf_add(accumulator, (Api.arr(h, d, f), word_at(h, d, f, at)))),
593      Equal.cong(Nat, F32 & Array<F32>,
594        n => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
595          Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(n)))),
596        Reduce.single_offset(rest, others, inner, offset), at, ReduceParallel.offset_shift(rest, others, inner, offset)),
597      Equal.cong(Array<F32> & F32, F32 & Array<F32>,
598        p => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
599          Reduce.leaf_add(accumulator, p)),
600        Array.get(F32, Api.arr(h, d, f), U32.from_nat(at)), (Api.arr(h, d, f), word_at(h, d, f, at)),
601        got(h, d, f, U32.from_nat(at)))))