B Getting started
If you want to just see it run — train a real neural network, end to end, with the Lean \(\to \) MLIR \(\to \) IREE pipeline from Part 1 — this appendix is the smallest path from zero to a trained model. Four tracks are included:
No-GPU Docker demo. Train MNIST on CPU in \(\sim \)5 minutes. Zero setup beyond Docker. Good if you don’t have a GPU and just want to confirm the pipeline works.
Native install (CUDA or ROCm). Train real image models (ResNet, MobileNet, EfficientNet, ViT) on your GPU. Requires an IREE build step but gives you the full book’s training runs.
Lean-only (no IREE). Read, build, and verify the proofs without running any neural-network training. Covers everything in Part 1 at the proof level.
One-command demo tiers. With the native setup (Track 2) in place, a single lake run mnist, cifar, or imagenette builds and runs a curated, time-budgeted group of verified trainers — no env vars or trainer list to assemble by hand.
Track 1: No-GPU Docker demo
The repository ships a Dockerfile that builds a CPU-only MNIST demo in a multi-stage image. Pulls Lean 4, builds IREE’s CPU runtime, compiles the MLP trainer, and bakes MNIST into the image. Once built, the final image is \(\sim \)300 MB and doesn’t require a GPU or Python.
git clone https://github.com/brettkoonce/lean4-mlir.git cd lean4-mlir docker build -t lean4-mlir-demo . docker run --rm lean4-mlir-demo
Output: the same three-layer MLP shown in Chapter 2 (784 \(\to \) 512 \(\to \) 512 \(\to \) 10) trains for 12 epochs on real MNIST, converges to \(\sim \)97.9% test accuracy. First build takes \(\sim \)10 minutes (dominated by IREE CMake + Ninja); subsequent docker run invocations reuse the cached image.
This is the fastest way to verify the Lean-to-MLIR-to-IREE pipeline actually works on your hardware before committing to a larger setup.
Track 2: Native install (CUDA or ROCm)
For training the larger networks (ResNet-34, MobileNet v2/v3/v4, EfficientNet, ViT) you need a GPU-accelerated IREE runtime. Steps:
Install Lean 4. Uses elan (like rustup for Rust) to manage toolchain versions.
curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh \ -sSf | sh
Build the IREE runtime for your GPU backend. See IREE_BUILD.md in the repo for the full CMake invocation. Once built, the FFI shim in ffi/ links against libiree_runtime_unified.a from your build tree.
Fetch data for the tier(s) you want to run:
./download_mnist.sh # MNIST (mnist tier) ./download_cifar.sh # CIFAR-10 (cifar tier) ./download_imagenette.sh # Imagenette 320px, ~300 MB (imagenette tier)
Build a trainer.
lake build resnet34-verified-adam
The targets are the verified trainers the three demo tiers (Track 4) bundle — build any one on its own:
lake run mnist: mnist-linear-verified, mnist-mlp-verified, mnist-cnn-verified
lake run cifar: cifar8-verified and cifar8-bn-verified, each in plain-SGD / -momentum / -adam (six in all)
lake run imagenette: resnet34-verified-adam, mobilenetv2-verified-adam, efficientnet-verified-adam, convnext-verified-adam, vit-verified-adam
(Unverified <arch>-train targets — vgg-train, resnet50-train, mobilenet-v3-train, efficientnet-v2-train, etc. — also build, for architectures outside the verified set.)
Run, setting the backend env vars. The FFI picks up IREE_BACKEND and IREE_CHIP at runtime.
NVIDIA / CUDA:
CUDA_VISIBLE_DEVICES=0 IREE_BACKEND=cuda IREE_CHIP=sm_86 \ .lake/build/bin/resnet34-verified-adam
(sm_86 is Ampere — RTX 30-series, A100. Use sm_89 for Ada/RTX 40-series, sm_90 for Hopper/H100.)
AMD / ROCm:
HIP_VISIBLE_DEVICES=0 IREE_BACKEND=rocm IREE_CHIP=gfx1100 \ .lake/build/bin/resnet34-verified-adam
(gfx1100 is the 7900 XT/XTX. Use gfx90a for the MI200 series, gfx942 for MI300.)
CPU fallback:
IREE_BACKEND=llvm-cpu .lake/build/bin/resnet34-verified-adam
(No chip argument. Slower than GPU but useful for debugging trainers without a working GPU setup.)
First-run compile is slow. IREE spends 10–15 minutes compiling a ResNet-sized training step to vmfb. Subsequent runs reuse the cached vmfb under .lake/build/ unless you delete it. The first-compile is a one-time cost per architecture, per GPU target.
Shell wrapper. The repo ships run.sh which sets the right env vars and tees output to a log file. Usage: ./run.sh <trainer> [gpu] [backend] (e.g. ./run.sh resnet34; ./run.sh efficientnet-v2 1 cuda).
Track 3: Proofs only (no IREE needed)
If you just want to read the book and verify the proofs, you can skip IREE entirely. The proofs don’t depend on the runtime.
git clone https://github.com/brettkoonce/lean4-mlir.git cd lean4-mlir lake exe cache get # pull ~5 GB of precompiled Mathlib lake build ProofsMinimal # the smallest end-to-end story (~seconds) lake build Proofs # type-checks the ENTIRE VJP proof suite
Start with ProofsMinimal. The full Proofs build type-checks every chapter at once — the real verification, but a lot to meet head-on. lake build ProofsMinimal instead compiles just the suite’s “hello world”: the linear softmax classifier — the book’s first end-to-end-verified example — in both halves of its claim.
LinearFaithfulPoC (\(\sim \)145 lines) — the faithfulness half: the StableHLO training step the codegen emits is, line for line, the certified softmax–cross-entropy gradient.
SgdDescentLinear (\(\sim \)255 lines) — the descent half: one SGD step on that certified gradient provably decreases the loss.
Their transitive closure is the minimum working set: those two files plus the shared StableHLO / tensor / float-bridge foundation, and nothing per-network beyond the linear case. This is the jumping-off point for the code. Open LeanMlir/Proofs/Foundation/LinearFaithfulPoC.lean and SgdDescentLinear.lean and read them top to bottom: you will have seen the entire shape of the argument the rest of the suite scales up — every later chapter is these same two halves at a larger network. The Start here section of LeanMlir/Proofs/README.md lays out that read order and flags which large files to skip on a first pass. Once ProofsMinimal builds green, graduate to the whole suite:
The Proofs target’s roots are the apex VJP modules; their transitive imports cover every chapter’s proof file, so this single build type-checks all of it against Lean’s kernel — that is the entire verification, no module list to keep current. Because the proof suite imports only Mathlib — never the codegen — Proofs skips the IREE/StableHLO layer entirely; it is also the default target, so a bare lake build runs the same check. (lake build LeanMlir builds the whole repo instead — codegen, trainers, and proofs.) Every VJP formula is proved against Mathlib’s \(\operatorname {fderiv}\) (see Appendix C).
To also run the numerical gradient checks:
python3 LeanMlir/Proofs/check_jacobians.py
Every closed-form Jacobian formula in the suite has a finite-difference test that compares the formula against the numerical derivative. The check typically completes in under a minute and reports pass/fail per formula.
Track 4: One-command demo tiers
Track 2 is the full manual path — pick a trainer, set the backend env vars, run. Once that setup works, you rarely want to assemble a trainer list and environment by hand every time. Three Lake scripts bundle the common cases into one command that builds and runs a curated, time-budgeted group of verified trainers back to back:
lake run mnist # verified MNIST: linear, MLP, CNN (~30 min) lake run cifar # ch.5 cifar8: SGD/momentum/Adam x bn/no-bn (~1 hr) lake run imagenette # the 5 Part-I nets at 224^2, 80-epoch AdamW (~37 h)
Each script streams every trainer live and tees it to <name>.log. Two conveniences make it genuinely one-click on top of the Track 2 install:
Backend auto-detects. cuda when an NVIDIA GPU is visible (nvidia-smi -L succeeds), otherwise rocm. Set IREE_BACKEND to override and LEAN_DEMO_GPU to pick the card (default 0).
The venv is found for you. The trainers shell out to iree-compile; the script prepends the project’s .venv/bin to PATH for the run, so it resolves without activating the venv first — the usual one-click footgun.
The tiers escalate by time budget. lake run mnist (\(\sim \)30 min) is the quickest end-to-end confirmation that the GPU path works; lake run cifar (\(\sim \)1 hr) is the Chapter 4 optimizer ablation in full — the six runs (SGD/momentum/AdamW \(\times \) BN/no-BN) behind that chapter’s accuracy curves and table; and lake run imagenette (\(\sim \)37 h on a single 7900 XTX) trains all five Part-I architectures (ResNet-34, MobileNetV2, EfficientNet-B0, ConvNeXt-T, ViT-Ti) at full \(224\times 224\) — a deliberate investment, not a quick look. Anything beyond these stays a single, deliberate run.sh invocation (Track 2).
Which tier fits your GPU? Before committing to one — especially the 37-hour Imagenette tier — run lake run benchmark. It briefly times a couple of the verified nets on your card and prints how long each chapter and each tier above would take on it:
$ lake run benchmark backend: rocm gpu: 0 ... 6 ResNet-34 9.5h 10 ViT 7.8h Full Part-1 training 42.8h lake run mnist 5m | cifar 34m | imagenette 42.2h
Backend auto-detects (CUDA or ROCm), so the numbers come out scaled to whatever card you run it on.
Common troubleshooting
HAL device not found on launch — usually means the IREE runtime wasn’t built with the requested backend’s HAL driver. CPU fallback (IREE_BACKEND=llvm-cpu) is the safest diagnostic: if it works on CPU but not GPU, the problem is in your IREE build’s HAL drivers, not in the Lean code.
lake build fails with Mathlib cache errors — run lake exe cache get to populate the local cache. Without it, the first build compiles all of Mathlib from source (\(\sim \)45 min). With it, Mathlib oleans are downloaded directly (\(\sim \)30 seconds).
First IREE compile appears stuck — it isn’t. ResNet-34’s full training-step compile takes 10–15 minutes. Watch RAM usage; if it’s climbing steadily the compile is progressing. If you want a faster check, try the MNIST MLP first (compiles in \(\sim \)30 seconds).