def kept_row

bendtorch-core@0.1.0.0 · src/autograd/reduce_fiber_proof.bend:790

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

790# A kept row walks the one digit the target names there.
791def kept_row(+size: Nat, +stride: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +extent: Nat, +others: List<&2, Nat>,
792  +target: Nat, +base: Nat, +inside: {Nat.is_lt(target, Nat.mul(extent, size)) == True{} : Bool},
793  +walk: RestWalk(rest, others, Nat.mod(target, size), stride), +bound: BoundBy<rest, size, stride>) ->
794  {masked(Nat.mul(extent, stride), Shape.Row{stride, size, right} <> rest, target, base, 0n) ==
795    visited(Shape.Row{stride, size, right} <> rest, extent <> others, Reduce.span(size, extent), target, base, Nat.div(target, size)) :
796    List<&2, Nat>}:
797  match size:
798    case 0n:
799      Empty.absurd({masked(Nat.mul(extent, stride), Shape.Row{stride, 0n, right} <> rest, target, base, 0n) ==
800        visited(Shape.Row{stride, 0n, right} <> rest, extent <> others, Reduce.span(0n, extent), target, base, Nat.div(target, 0n)) :
801        List<&2, Nat>},
802        Arith.lt_zero(target,
803          %Nats.mul_zero_right(extent) : {Nat.is_lt(target, _) == True{} : Bool}
804          inside))
805    case 1n+l:
806      +dt = Nat.div(target, 1n+l)
807      +tr = Nat.mod(target, 1n+l)
808      +w = visited(rest, others, Reduce.first_span(rest, others), tr, Nat.add(base, Nat.mul(right, dt)), Reduce.first_digit(rest, tr))
809      %Equal.sym(List<&2, Nat>, masked(Nat.mul(extent, stride), Shape.Row{stride, 1n+l, right} <> rest, target, base, 0n),
810        masked(stride, rest, tr, Nat.add(base, Nat.mul(right, dt)), 0n),
811        kept_hit(extent, stride, 1n+l, right, rest, target, base, 0n, Arith.le_zero(dt),
812          div_below(target, extent, 1n+l, inside), decompose(target, 1n+l, {==}), bound)) :
813        {_ == List.append(&2, Nat, w, Nil{}) : List<&2, Nat>}
814      %Equal.sym(List<&2, Nat>, masked(stride, rest, tr, Nat.add(base, Nat.mul(right, dt)), 0n), w,
815        rest_at(rest, others, tr, stride, walk, Nat.add(base, Nat.mul(right, dt)))) :
816        {_ == List.append(&2, Nat, w, Nil{}) : List<&2, Nat>}
817      Equal.sym(List<&2, Nat>, List.append(&2, Nat, w, Nil{}), w, append_nil(w))