def fiber_step

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

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

430# One digit of a row whose rows below take one digit adds the word at its
431# strided offset.
432def fiber_step(+stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +extent: Nat, +others: List<&2, Nat>,
433  +single: {Reduce.single_rows(rest, others) == True{} : Bool}, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat, +accumulator: F32,
434  -array: Array<F32>, -kept: Array<F32>, -word: F32,
435  -read: {(kept, word) == Array.get(F32, array, U32.from_nat(Reduce.strided_offset(right, base, digit, Reduce.single_offset(rest, others, Nat.mod(target, left), 0n)))) : Array<F32> & F32}) ->
436  Generic.Fact<{Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 1n+fuel, target, base, digit, accumulator, array) ==
437    Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, F32.add(accumulator, word), kept) : F32 & Array<F32>}>:
438  +inner = Nat.mod(target, left)
439  +lift = Reduce.single_offset(rest, others, inner, 0n)
440  +offset = Nat.add(base, Nat.mul(right, digit))
441  +at = Reduce.strided_offset(right, base, digit, lift)
442  Generic.Fact{Equal.trans(F32 & Array<F32>,
443    Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
444      Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)),
445    Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
446      Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset))))),
447    Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
448      Reduce.leaf_add(accumulator, (kept, word))),
449    Equal.cong(F32 & Array<F32>, F32 & Array<F32>,
450      p => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k), p),
451      Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array),
452      Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))),
453      {\{Generic.Fact: p => p} : Generic.Fact<({Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) ==
454        Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))) : F32 & Array<F32>})> ->
455        {Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) ==
456          Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))) : F32 & Array<F32>}}(
457        single_fiber(rest, others, single, inner, offset, accumulator, array))),
458    Equal.trans(F32 & Array<F32>,
459      Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
460        Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset))))),
461      Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
462        Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(at)))),
463      Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
464        Reduce.leaf_add(accumulator, (kept, word))),
465      Equal.cong(Nat, F32 & Array<F32>,
466        n => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
467          Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(n)))),
468        Reduce.single_offset(rest, others, inner, offset), at, offset_shift(rest, others, inner, offset)),
469      Equal.cong(Array<F32> & F32, F32 & Array<F32>,
470        p => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, fuel, target, base, 1n+digit, v, k),
471          Reduce.leaf_add(accumulator, p)),
472        Array.get(F32, array, U32.from_nat(at)), (kept, word),
473        Equal.sym(Array<F32> & F32, (kept, word), Array.get(F32, array, U32.from_nat(at)), read))))}