def shifted

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith

474# On a row of stride s, left stride L and right stride R, the positions
475# d * s + q .. d * s + q + n - 1 have digit d and walk the rows below at
476# q .. q + n - 1.
477def shifted(+n: Nat, +stride: Nat, +left: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +target: Nat, +inner: Nat, +base: Nat,
478  +digit: Nat, +q: Nat, +room: {Nat.is_le(Nat.add(q, n), stride) == True{} : Bool}, +mode: MaskMode<left, digit, target, inner>) ->
479  {masked(n, Shape.Row{stride, left, right} <> rest, target, base, Nat.add(Nat.mul(digit, stride), q)) ==
480    masked(n, rest, inner, Nat.add(base, Nat.mul(right, digit)), q) : List<&2, Nat>}:
481  match n:
482    case 0n:
483      {==}
484    case 1n+m:
485      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
486      +c = Nat.add(Nat.mul(digit, stride), q)
487      +below = lt_of_room(q, m, stride, room)
488      +dq = div_of(c, stride, digit, q, below, {==})
489      +mq = mod_of(c, stride, digit, q, below, {==})
490      +wl = Shape.walk_left(rest, q)
491      +wr = Shape.walk_right(rest, q)
492      +tail = masked(m, rows, target, base, 1n+c)
493      +goal_right = kept(Nat.is_eq(wl, inner), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr),
494        masked(m, rest, inner, Nat.add(base, Nat.mul(right, digit)), 1n+q))
495      %Equal.sym(Nat, Nat.div(c, stride), digit, dq) :
496        {kept(Nat.is_eq(Nat.add(Nat.mul(left, _), Shape.walk_left(rest, Nat.mod(c, stride))), target),
497          Nat.add(base, Nat.add(Nat.mul(right, _), Shape.walk_right(rest, Nat.mod(c, stride)))), tail) == goal_right : List<&2, Nat>}
498      %Equal.sym(Nat, Nat.mod(c, stride), q, mq) :
499        {kept(Nat.is_eq(Nat.add(Nat.mul(left, digit), Shape.walk_left(rest, _)), target),
500          Nat.add(base, Nat.add(Nat.mul(right, digit), Shape.walk_right(rest, _))), tail) == goal_right : List<&2, Nat>}
501      %mode_mask(left, digit, target, inner, mode, wl) :
502        {kept(_, Nat.add(base, Nat.add(Nat.mul(right, digit), wr)), tail) == goal_right : List<&2, Nat>}
503      %Equal.sym(Nat, Nat.add(base, Nat.add(Nat.mul(right, digit), wr)), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr),
504        Index.add_assoc(base, Nat.mul(right, digit), wr)) :
505        {kept(Nat.is_eq(wl, inner), _, tail) == goal_right : List<&2, Nat>}
506      %Arith.add_succ(Nat.mul(digit, stride), q) :
507        {kept(Nat.is_eq(wl, inner), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr), masked(m, rows, target, base, _)) == goal_right : List<&2, Nat>}
508      Equal.cong(List<&2, Nat>, List<&2, Nat>, xs => kept(Nat.is_eq(wl, inner), Nat.add(Nat.add(base, Nat.mul(right, digit)), wr), xs),
509        masked(m, rows, target, base, Nat.add(Nat.mul(digit, stride), 1n+q)),
510        masked(m, rest, inner, Nat.add(base, Nat.mul(right, digit)), 1n+q),
511        shifted(m, stride, left, right, rest, target, inner, base, digit, 1n+q, le_step(q, m, stride, room), mode))