def single_arr

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
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

525def single_arr(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +single: {Reduce.single_rows(rows, extents) == True{} : Bool},
526  +target: Nat, +base: Nat, +accumulator: F32, +h: List<&2, Api.Effect>, +d: Nat, +f: F32) ->
527  {Reduce.fiber(rows, extents, Reduce.first_span(rows, extents), target, base, Reduce.first_digit(rows, target), accumulator, Api.arr(h, d, f)) ==
528    Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rows, extents, target, base)))) : F32 & Array<F32>}:
529  match rows extents:
530    case Nil{} +extents:
531      {==}
532    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others:
533      +one = ReduceParallel.span_one(left, extent, Arith.and_left(Nat.is_eq(Reduce.span(left, extent), 1n), Reduce.single_rows(rest, others), single))
534      +inner = Nat.mod(target, left)
535      +offset = Nat.add(base, Nat.mul(right, Nat.div(target, left)))
536      %Equal.sym(Nat, Reduce.span(left, extent), 1n, one) :
537        {Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, _, target, base, Nat.div(target, left), accumulator, Api.arr(h, d, f)) ==
538          Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))) : F32 & Array<F32>}
539      Equal.trans(F32 & Array<F32>,
540        Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k),
541          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f))),
542        Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k),
543          Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset))))),
544        Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))),
545        Equal.cong(F32 & Array<F32>, F32 & Array<F32>,
546          p => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k), p),
547          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)),
548          Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))),
549          single_arr(rest, others, Arith.and_right(Nat.is_eq(Reduce.span(left, extent), 1n), Reduce.single_rows(rest, others), single),
550            inner, offset, accumulator, h, d, f)),
551        ReduceParallel.leaf_kept(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k),
552          value => kept => {==}, accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))))
553    case Shape.Row{+stride, +left, +right} <> +rest Nil{}:
554      Empty.absurd({Reduce.fiber(Shape.Row{stride, left, right} <> rest, Nil{}, Reduce.first_span(Shape.Row{stride, left, right} <> rest, Nil{}), target, base,
555        Reduce.first_digit(Shape.Row{stride, left, right} <> rest, target), accumulator, Api.arr(h, d, f)) ==
556        Reduce.leaf_add(accumulator, Array.get(F32, Api.arr(h, d, f), U32.from_nat(Reduce.single_offset(Shape.Row{stride, left, right} <> rest, Nil{}, target, base)))) :
557        F32 & Array<F32>}, Arith.clash(single))