def mask_at

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

import Base
import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats

194# On a kept row of left stride `size`, the position with digit d and
195# walk w below size lands on t exactly when d is t's quotient and w its
196# remainder.
197def mask_at(+size: Nat, +d: Nat, +w: Nat, +t: Nat, given: Decomposed<t, size>, same: {d == Nat.div(t, size) : Nat}) ->
198  {Nat.is_eq(Nat.add(Nat.mul(size, d), w), t) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}:
199  Decomposed{below, sum} = given
200  %sum :
201    {Nat.is_eq(Nat.add(Nat.mul(size, d), w), _) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}
202  %same : {Nat.is_eq(Nat.add(Nat.mul(size, d), w), Nat.add(Nat.mul(_, size), Nat.mod(t, size))) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}
203  %Nats.mul_comm(size, d) : {Nat.is_eq(Nat.add(Nat.mul(size, d), w), Nat.add(_, Nat.mod(t, size))) == Nat.is_eq(w, Nat.mod(t, size)) : Bool}
204  eq_add_cancel(Nat.mul(size, d), w, Nat.mod(t, size))