Documentation

LeanMlir.GradcheckHelpers

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 ViT gradcheck tests (tests/TestSDPA.lean, TestMHSA, TestViTBlock, TestViTTiny). All Lean4.

source
def ViTGradcheck.runDevice :
IO String

The iree-run-module device for the IREE_BACKEND the .vmfb was compiled for (default cuda, as in ireeCompileArgs): rocm runs on hip, llvm-cpu on local-task.

Equations
Instances For
    source
    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.
    Instances For
      source
      def ViTGradcheck.dot (a b : Array Float) :
      Float
      Equations
      Instances For
        source
        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

        Adjoint/finite-difference gradcheck of a compiled fwd/back pair, checking ⟨back(in, dOut), v⟩ ≈ (Φ(+ε) − Φ(−ε))/2ε with Φ(s) = ⟨fwd(in + s·v), dOut⟩ for random v. inShapes/inLens describe the perturbed inputs (in arg order), outShape/outLen the forward's single output. fixed inputs (concrete (shape,values)) are passed to BOTH fwd and back, never perturbed, and have no expected gradient — e.g. a ViT input image. Forward arg order is fixed ++ params; the backward takes fixed ++ params ++ dOut and returns one grad per PARAM (in order). 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.
        Instances For
          source
          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

          adjointGradcheckFixed with no fixed inputs: the backward returns one gradient per input.

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