PROOF.bend source

PROOF.bend on the hub · documented module

# PROOF.bend -- proofs of the laws in LAWS.bend. The gate:#   bend PROOF.bend   =>  "All terms check."## Proof notes (learned on Bend 2.0.4):#   - induction is a recursive call on a structurally smaller pattern variable#   - `%e : P` rewrites the goal: P is the goal with `_` marking the b-side of#     e's equation, and the goal becomes P with `a` in that place#   - evidence of `Empty` (clashes like 2n-vs-3n broadcast) is discharged by a#     match with no casesimport Baseimport ./tinygrad/shape.bend as Shimport ./LAWS.bend as Laws# ---------- LAW 1: broadcast commutes ----------def max_comm(+a: Nat, +b: Nat) -> {Sh.maxn(a, b) == Sh.maxn(b, a) : Nat}:  match a b:    case 0n, 0n:      {==}    case 0n, 1n+b2:      {==}    case 1n+a2, 0n:      {==}    case 1n+a2, 1n+b2:      %max_comm(a2, b2) : {1n+Sh.maxn(a2, b2) == 1n+_ : Nat}      {==}def Laws.bcast_commute(s, t):  match s t:    case Nil{}, Nil{}:      {==}    case Nil{}, Con{b, bs}:      {==}    case Con{a, as}, Nil{}:      {==}    case Con{a, as}, Con{b, bs}:      %max_comm(a, b) : {Sh.maxn(a, b) <> Sh.bcast(as, bs) == _ <> Sh.bcast(bs, as) : List<&2, Nat>}      %Laws.bcast_commute(as, bs) : {Sh.maxn(a, b) <> Sh.bcast(as, bs) == Sh.maxn(a, b) <> _ : List<&2, Nat>}      {==}# ---------- LAW 2: matmul shape ----------def Laws.matmul_shape(m, k, n):  {==}# ---------- LAW 3: broadcast numel ----------# a dominating dim yields the maxdef max_bdim2(d2: Nat, e2: Nat, ev: Sh.bdim2(d2, e2)) -> {Sh.maxn(d2, e2) == e2 : Nat}:  match d2 e2:    case 0n, 0n:      {==}    case 0n, 1n+e3:      {==}    case 1n+d3, 0n:      match ev:    case 1n+d3, 1n+e3:      %max_bdim2(d3, e3, ev) : {1n+Sh.maxn(d3, e3) == 1n+_ : Nat}      {==}def max_bdim(d: Nat, e: Nat, ev: Sh.bdim(d, e)) -> {Sh.maxn(d, e) == e : Nat}:  match d e:    case 0n, _:      match ev:    case 1n+d2, 0n:      match ev:    case 1n+d2, 1n+e2:      Equal.cong(Nat, Nat, x => 1n+x, Sh.maxn(d2, e2), e2, max_bdim2(d2, e2, ev))def Laws.bcast_numel(s, t, c):  match s t:    case Nil{}, Nil{}:      {==}    case Nil{}, Con{b, bs}:      {==}    case Con{_, _}, Nil{}:      c    case Con{a, as}, Con{b, bs}:      (dev, rest) = c      ih = Laws.bcast_numel(as, bs, rest)      c1 = Equal.cong(Nat, Nat, x => (Nat.mul(Sh.maxn(a, b), x) : Nat), Sh.numel(Sh.bcast(as, bs)), Sh.numel(bs), ih)      c2 = Equal.cong(Nat, Nat, x => (Nat.mul(x, Sh.numel(bs)) : Nat), Sh.maxn(a, b), b, max_bdim(a, b, dev))      Equal.trans(Nat, Nat.mul(Sh.maxn(a, b), Sh.numel(Sh.bcast(as, bs))), Nat.mul(Sh.maxn(a, b), Sh.numel(bs)), Nat.mul(b, Sh.numel(bs)), c1, c2)# ---------- LAW 4: gradient accumulation is lossless ----------def Laws.grad_accum_len(xs, ys):  match xs:    case Nil{}:      {==}    case Con{h, t}:      Equal.cong(Nat, Nat, x => 1n+x,        List.length(&2, Nat, List.append(&2, Nat, t, ys)),        Nat.add(List.length(&2, Nat, t), List.length(&2, Nat, ys)),        Laws.grad_accum_len(t, ys))def sub_diag(+a: Nat) -> {Nat.sub(a, a) == 0n : Nat}:  match a:    case 0n:      {==}    case 1n+p:      %sub_diag(p) : {Nat.sub(p, p) == _ : Nat}      {==}def add_zero_right(+a: Nat) -> {Nat.add(a, 0n) == a : Nat}:  match a:    case 0n:      {==}    case 1n+p:      %add_zero_right(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}      {==}# ---------- LAW 6: reduce count (axis 0) ----------def Laws.reduce_numel(s):  match s:    case Nil{}:      {==}    case Con{h, t}:      {==}