Verified Deep Learning with Lean 4

Getting started

Install first, read while it downloads. Four steps put a trained network on your GPU, and after that every chapter opens with a command you can run. This is the default path, XLA through PJRT, and the one every number in this book was measured on. There is no compiler to build. The proofs-only track, the six GPU platforms, the second lowerer and troubleshooting are Appendix A; ImageNet-1k is Appendix C.

Install

  1. 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

  2. Get an XLA PJRT plugin, one prebuilt shared library. It ships under a JAX name because that is the only place XLA’s GPU plugins are packaged, but these wheels carry the plugin and nothing else (not jax, not jaxlib), and the shim opens the .so directly through the PJRT C API, so no interpreter runs at training time. 

    python3 -m venv .venv && . .venv/bin/activate
pip install jax-cuda12-pjrt     # NVIDIA, CUDA 12 -- the PJRT plugin ONLY
pip install jax-rocm7-pjrt      # AMD            -- no jax, no jaxlib

    The shim looks under .venv/lib/python3.12/site-packages/jax_plugins/, which is where that command puts it on a Python 3.12 venv. On any other Python, or for a plugin anywhere else (a CUDA 13 card’s jax-cuda13-pjrt, for one), $PJRT_PLUGIN names the .so and always wins. The plugin is version-matched to your CUDA/ROCm install: a 500 MB download, not a build.

  3. Build the FFI shim (needs only libc and dlopen): 

    gcc -fPIC -O2 -shared ffi/pjrt_ffi.c -ldl -o ffi/libpjrt_ffi.so

  4. Fetch data for the tier(s) you want: 

    ./scripts/datasets/download_mnist.sh       # MNIST       (mnist tier)
./scripts/datasets/download_cifar.sh       # CIFAR-10    (cifar tier)
./scripts/datasets/download_imagenette.sh  # Imagenette 320px, 330 MB (imagenette)

  5. Build and run a trainer. No backend selection: XLA finds the device. The linear MNIST classifier of Chapter 1 trains in seconds: 

    lake build mnist-linear-verified
.lake/build/bin/mnist-linear-verified data

    The trainers take the data root as their one argument and default to data. The Imagenette nets build the same way (lake build resnet34-verified-adam) and run for the better part of an hour each; the tiers below bundle them.

One command per tier

With the install in place, one command builds and runs a time-budgeted group of verified trainers. 

lake run mnist        # verified MNIST: linear, MLP, CNN              (~1 min)
lake run cifar        # ch.4 cifar8w: SGD/momentum/AdamW x 3 norms   (~13 min)
lake run imagenette   # the Part-I nets at 224^2, 80-epoch AdamW      (~6 h)
lake run imagenet     # ImageNet-1k on 4 GPUs: plans; `start` runs it (Appendix C)

The three times are XLA wall-clocks on one RTX 4060 Ti, resident; the fourth tier is weeks of four-card time and has Appendix C to itself. The Imagenette figure is the five chapter nets at their chapters’ times, about 3.7 hours between them, plus ResNet-50 (about an hour and a half) and MobileNetV4 (under an hour), which ride in the same tier. Each trainer’s transcript is teed to runs/<date>-<trainer>/. Each tier also has an -iree twin, which is Appendix A’s business. The tiers bundle these targets, and any one builds on its own:

  • mnist: mnist-linear-verified, mnist-mlp-verified, mnist-cnn-verified

  • cifar: cifar8w-ablation, cifar8w-bn-ablation and cifar8wb-bbn-ablation — no BatchNorm, per-image and batch statistics — each running SGD / momentum / AdamW in sequence, the last in fp32 and bf16: twelve arms in three binaries

  • imagenette, in chapter order, each side quest after the chapter it belongs to: resnet34-verified-adam, resnet50-verified-adam, mobilenetv2-verified-adam, mobilenetv4-verified-adam, efficientnet-verified-adam, convnext-verified-adam, vit-verified-adam