def operand_step

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul_PROOF.bend:458

import Base
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/batched_matmul_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
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

458# Term k + 1 of an operand sits one stride after term k.
459def operand_step(+order: Matrix.Order, +inner: Nat, +outers: Nat, +e1: Nat, +e2: Nat, +b1: Nat, +b2: Nat, +outer: Nat, +k: Nat) ->
460  {Nat.add(Matrix.stride(order, outers), Laws.operand_offset(order, inner, outers, e1, e2, b1, b2, outer, k)) ==
461    Laws.operand_offset(order, inner, outers, e1, e2, b1, b2, outer, 1n+k) : Nat}:
462  match order:
463    case Matrix.Kept{}:
464      %Equal.sym(Nat, Nat.add(Nat.mul(outer, inner), 1n+k), 1n+Nat.add(Nat.mul(outer, inner), k), Arith.add_succ(Nat.mul(outer, inner), k)) :
465        {1n+Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k)) == Nat.add(batch_base(inner, outers, e1, e2, b1, b2), _) : Nat}
466      %Equal.sym(Nat, Nat.add(batch_base(inner, outers, e1, e2, b1, b2), 1n+Nat.add(Nat.mul(outer, inner), k)),
467        1n+Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k)),
468        Arith.add_succ(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k))) :
469        {1n+Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k)) == _ : Nat}
470      {==}
471    case Matrix.Swapped{}:
472      %Index.add_assoc(outers, Nat.mul(k, outers), outer) :
473        {Nat.add(outers, Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(k, outers), outer))) == Nat.add(batch_base(inner, outers, e1, e2, b1, b2), _) : Nat}
474      Nats.add_swap(outers, batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(k, outers), outer))