Lean4 numerical gradcheck harness (no numpy)

Shells out to iree-run-module to execute compiled @*_fwd/@*_back .vmfb, then runs the adjoint / finite-difference dot-product test: for a forward f with VJP J·ᵀ, the backward gives g_i = (Jᵀ dOut)_i, and for random perturbation directions v_i, Σ_i ⟨g_i, v_i⟩ = ⟨Jᵀ dOut, v⟩ = ⟨dOut, J v⟩ = (Φ(+ε) − Φ(−ε)) / 2ε, where Φ(s) := ⟨f(inputs + s·v), dOut⟩. One backward run + two forward runs validate ALL input gradients at once — catching transpose/axis bugs that iree-compile (type-checking only) cannot.

Used by the ch10 ViT de-risk tests (TestSDPA/TestMHSA/TestViTBlock). All Lean4.

def ViTGradcheck.pow10 (k : Int) : Float

10^k as a Float (k may be negative).

Equations

• One or more equations did not get rendered due to their size.

def ViTGradcheck.digitsToNat (cs : List Char) : Nat

Equations

• ViTGradcheck.digitsToNat cs = List.foldl (fun (acc : Nat) (c : Char) => acc * 10 + (c.toNat - '0'.toNat)) 0 cs

def ViTGradcheck.splitAtChar (c : Char) (xs : List Char) : List Char × Option (List Char)

Split a List Char at the first occurrence of c (the separator dropped).

Equations

• ViTGradcheck.splitAtChar c xs = match List.span (fun (x : Char) => x != c) xs with | (pre, []) => (pre, none) | (pre, head :: post) => (pre, some post)

def ViTGradcheck.parseFloat (tok : String) : Float

Parse one iree-printed float token (-0.00623606, 1.3e-05, 42), entirely over List Char (robust to the String/Slice API churn).

Equations

• One or more equations did not get rendered due to their size.

def ViTGradcheck.parseResults (out : String) : Array (Array Float)

Extract the parsed result buffers (in result[i] order) from an iree-run-module stdout: each value line is …xf32=[a b][c d]….

Equations

• One or more equations did not get rendered due to their size.

def ViTGradcheck.runFn (vmfb fn : String) (inputs : List (String × Array Float)) : IO (Array (Array Float))

Run a compiled .vmfb function; inputs are (shapeStr, flatValues).

Equations

• One or more equations did not get rendered due to their size.

def ViTGradcheck.randVec (seed n : Nat) : Array Float

Deterministic LCG pseudo-random Array Float in [-1,1], length n.

Equations

• One or more equations did not get rendered due to their size.

def ViTGradcheck.dot (a b : Array Float) : Float

Equations

• ViTGradcheck.dot a b = Array.foldl (fun (acc : Float) (x : Float × Float) => match x with | (x, y) => acc + x * y) 0.0 (a.zip b)

def ViTGradcheck.axpy (a : Float) (x y : Array Float) : Array Float

y + a·x (elementwise).

Equations

• ViTGradcheck.axpy a x y = Array.map (fun (x : Float × Float) => match x with | (yi, xi) => yi + a * xi) (y.zip x)

def ViTGradcheck.adjointGradcheck (label fwdVmfb fwdFn backVmfb backFn : String) (inShapes : List String) (inLens : List Nat) (outShape : String) (outLen : Nat) (seedBase : Nat := 0) (eps : Float := 10e-4) (tol : Float := 10e-3) : IO Bool

Adjoint/finite-difference gradcheck of a compiled fwd/back pair. inShapes/inLens describe the forward inputs (in arg order); the backward is expected to return one gradient per input in the same order. outShape/ outLen describe the forward's single output (the dOut cotangent). Returns true iff the relative error is below tol (default 1e-2, for f32).

Equations

• One or more equations did not get rendered due to their size.

def ViTGradcheck.adjointGradcheckFixed (label fwdVmfb fwdFn backVmfb backFn : String) (fixed : List (String × Array Float)) (inShapes : List String) (inLens : List Nat) (outShape : String) (outLen : Nat) (seedBase : Nat := 0) (eps : Float := 10e-4) (tol : Float := 10e-3) : IO Bool

Like adjointGradcheck but with fixed inputs (concrete (shape,values)) that are passed to BOTH fwd and back, never perturbed, and have no expected gradient — e.g. a ViT input image (first layer ⇒ no image grad). The forward arg order is fixed ++ params; the backward is fixed ++ params ++ dOut and returns one grad per PARAM (in order).

Equations

• One or more equations did not get rendered due to their size.