def split_step

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

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

47def split_step(+t: Nat, +size: Nat, +q: Nat, +r: Nat, below: {Nat.is_lt(r, size) == True{} : Bool},
48  sum: {Nat.add(Nat.mul(q, size), r) == t : Nat}, last: Shape.BoolCase<Nat.is_lt(1n+r, size)>) -> Split<1n+t, size>:
49  match last:
50    case Shape.BTrue{equal}:
51      Split{q, 1n+r, Equal.sym(Bool, True{}, Nat.is_lt(1n+r, size), equal),
52        Equal.trans(Nat, Nat.add(Nat.mul(q, size), 1n+r), 1n+Nat.add(Nat.mul(q, size), r), 1n+t,
53          Arith.add_succ(Nat.mul(q, size), r), Equal.cong(Nat, Nat, n => 1n+n, Nat.add(Nat.mul(q, size), r), t, sum))}
54    case Shape.BFalse{equal}:
55      +whole = lt_succ_eq(r, size, below, equal)
56      Split{1n+q, 0n,
57        %Equal.sym(Nat, size, 1n+r, whole) : {Nat.is_lt(0n, _) == True{} : Bool}
58        {==},
59        %Equal.sym(Nat, Nat.add(Nat.mul(1n+q, size), 0n), Nat.mul(1n+q, size), Broadcast.add_zero(Nat.mul(1n+q, size))) : {_ == 1n+t : Nat}
60        %Equal.sym(Nat, size, 1n+r, whole) : {Nat.add(_, Nat.mul(q, size)) == 1n+t : Nat}
61        %Equal.sym(Nat, Nat.add(r, Nat.mul(q, size)), Nat.add(Nat.mul(q, size), r), add_comm(r, Nat.mul(q, size))) : {1n+_ == 1n+t : Nat}
62        Equal.cong(Nat, Nat, n => 1n+n, Nat.add(Nat.mul(q, size), r), t, sum)}