Verified Deep Learning with Lean 4

A Toolchain and platforms

Getting started, at the front of the book, is the smallest path from zero to a trained model. This appendix is the rest of the toolchain: the proofs-only track, the six GPU platforms, troubleshooting, and the second lowerer.

Table A-1 is the scale. What you install depends on how far up it you want to go, and the steps are not evenly spaced.

Table A-1 The four tracks: what each lets you do and what it needs.

What you want

What you need

Track

Where

Read and check the proofs

Lean. No GPU.

1

below

MNIST / CIFAR / Imagenette

Lean + one plugin file.

2, 3

Getting started

Full ImageNet-1k

the above, plus Python + tfds.

4

Appendix C

Everything in Part 1 of this book is the middle row. Nothing runs through an interpreter at training time: you download one shared library and run a Lean binary. Python appears twice before that, as the pip that fetches the plugin wheel and in the Imagenette download script, which resizes with Pillow; the shim is one gcc line. Only the ImageNet tier streams its data through Python, and it has Appendix C to itself.

Track 1: Proofs only (no GPU needed)

Read, build, and check the proofs without running any training. This track needs neither lowerer: the theorems are about the emitted graph, not about anything that executes it. It comes at three levels, each asking you to trust less than the last.

Level 1: build a proof

lake exe cache get           # Mathlib oleans, ~30 s (vs ~45 min from source)
lake build ProofsMinimal     # two files: the faithfulness fold + SGD descent
lake build Proofs            # the proof suite
lake build Certs             # the certificates CI checks; subsumes Proofs (slow)

ProofsMinimal is the smallest thing that is still a real result: LinearFold (the rendered linear layer denotes the Mathlib fderiv math) and SgdDescent.Linear (the step decreases the loss). Start there if you want to see one theorem land before committing to the full build. The Lipschitz and IBP scorecards are a separate target, CertsHeavy.

Level 2: check the axioms

#print axioms on any theorem shows its transitive axiom closure. Every result in this book closes over Lean’s three core axioms and nothing else, and Appendix B is the accounting.

Level 3: re-check with an independent kernel

#print axioms and lake build both run through Lean’s elaborator, so they share a trust path with whatever the elaborator did.

tests/comparator/run.sh

re-runs Lean’s kernel typechecker, and five independent checkers beside it, over 110 theorems, sandboxed, with an axiom allowlist of exactly the three. It needs a built Certs and bwrap. Which 110, why thirteen of them are stated in a file that imports nothing but Mathlib, and what the run adds to #print axioms is §B.2.2 in Appendix B.

Platforms

The target is the middle row of Table A-1 — MNIST, CIFAR and Imagenette, Tracks 2 and 3 — on the six GPU platforms of Table A-2, each one plugin file and (off the Mac) one device variable away from the same Lean binaries, the plugin pinned per platform.

Table A-2 The GPU platforms, the plugin each is pinned to, and where each stands.

Platform

Plugin, pinned

Where it stands

NVIDIA, CUDA 12

jax-cuda12-pjrt 0.11.0 (CUDA 12.9, cuDNN 9.23)

the book’s numbers; Tracks 1–4

NVIDIA, CUDA 13

jax-cuda13-pjrt 0.11.1 (CUDA 13.3, cuDNN 9.24)

Track 4 on record: the four-card ImageNet runs

AMD, ROCm 7.2

jax-rocm7-pjrt 0.11.0 (the second pip line of Getting started)

in the code (HIP_VISIBLE_DEVICES, the suite’s –backend rocm)

Jetson Orin Nano 8 GB, JetPack 6.2

jax 0.11.1 built for sm_87 against cuDNN 9.12, per deploy/ORIN.md

Track 2 on record, and of Track 3 the MNIST CNN, the CIFAR tier and ViT-Tiny over 80 epochs; lake run imagenette-orin, the six-row Imagenette tier at the board’s memory, is the next run

Apple Silicon, Metal (M1 MacBook Air, 16 GB, macOS 26)

jax-graft 0.0.5 (metal-pjrt-core, platform mtl: jaxlib 0.11.2’s XLA with a Metal backend), per deploy/MAC.md

Track 2 and the MNIST tier of Track 3 on record in deploy/MAC.md (92.10 / 97.79 / 98.71%, against CUDA’s medians 92.10 / 97.78 / 98.78); CIFAR and ResNet-34 train, their full runs to come

Intel Arc Pro B50, oneAPI (16 GB, Battlemage, kernel 7.0)

jax-oneapi-pjrt 0.11.2 (Intel’s OpenXLA plugin, platform oneapi, its oneAPI runtime as wheels), per deploy/ARC_B50.md

Track 2 and the MNIST and CIFAR tiers of Track 3 on record, matching CUDA (92.10 / 97.83 / 98.66%; CIFAR 76.24%), bf16 included; all seven Imagenette nets’ 80 epochs on record (the slot at PCIe Gen3), each inside CUDA’s five-seed range or within two images of it

The CUDA 12 and ROCm pins are the lock files under jax/; the CUDA 13 box’s is scripts/jobs/_box.sh, the Orin’s its build script, the Mac’s and the Arc’s their setup scripts. Three of the six are there to make a point. The Orin is the core sequence on a commodity computer rather than a training box — the board, a plugin built for it and the tier’s memory recipe, nothing else — and deploy/ORIN.md is its runbook. The Mac is the same point on a laptop: Metal through a PJRT plugin, unified memory (nothing is staged to the device), float32 only (the MNIST renders use no float64), and deploy/MAC.md is its runbook. The Arc is the third vendor’s desktop card through the same seam, Intel’s PJRT plugin and nothing else from this side, with one shim workaround for a race in Intel’s SYCL runtime, and deploy/ARC_B50.md is its runbook. A new box is admitted by scripts/platform/check.sh: the plugin loads, the shim’s guards hold and the committed artifacts match XLA:CPU goldens, in minutes and without Lean. PLATFORMS.md is the support matrix, the last run of each tier on each box; where a cell is empty, as the Mac’s are, the runbook is the record.

Common troubleshooting

  • no lowerer shim could be loaded. No ffi/libpjrt_ffi.so. Build it with the gcc line in Getting started, which needs nothing but libc. (Under LEAN_MLIR_LOWERER=iree the same message says wanted: iree, and the shim it wants is libiree_ffi.so.)

  • $PJRT_PLUGIN is set but did not load, or no plugin found. The XLA plugin is missing or is built for a different CUDA/ROCm version than the one installed. Installing the plugin wheel into the repo’s own .venv on Python 3.12, as in Getting started, puts it where the shim looks by default; anywhere else, set $PJRT_PLUGIN.

  • Picking a card, and parameter residency. As with any GPU program, the card is chosen by CUDA_VISIBLE_DEVICES on NVIDIA, HIP_VISIBLE_DEVICES on AMD, and ZE_AFFINITY_MASK on the Arc. Parameters stay on the device between steps by default, and the tiers are timed that way. PJRT_FFI_RESIDENT=0 forces the copying path, which the residency gate uses as its control: the same run is 1.3–3\(\times \) slower and its numerics are unchanged.

  • 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), and with it oleans download in \(\sim \)30 seconds.

  • A run refuses to start over a missing shim (ImageNet tier). Run scripts/gen_shims.sh; Appendix C.

  • lake run <job> plan refuses with “OLDER than the render” (ImageNet tier). The job’s precheck compares the binary’s timestamp with its render’s, and a regenerated render trips it. lake build <exe> clears it; a bare lake run <job> builds first and never sees it.

The second lowerer

Everything above runs on XLA, which is the default and needs nothing from this section. There is a second trusted lowerer, IREE, and the reason to know it exists is not that you might train on it.

Both consume the same proven StableHLO and sit in the same trusted tier; neither is verified, and swapping them changes nothing about what is proved. XLA is faster here — \(2.3\times \) on the MNIST linear net, \(4.9\times \) on the CNN — because it dispatches to the vendor’s hand-tuned kernels, where IREE generates every kernel itself. That is also why IREE’s runtime is \(392\times \) smaller, \(1.3\) MB against the XLA plugin’s \(505\) MB. The size gap is not a trust gap: most of the plugin is vendor kernel libraries it calls, where IREE generates the equivalent, and both bodies of code are unverified. Choosing the small runtime relocates the trusted base rather than shrinking it.

What IREE is for here is the oracles. The differential checks of §B.4.2 lower the Lean side through IREE and compare against a JAX reference running through XLA, and it is the fact that those are two independent compilers — different code generation, different runtime, different kernels — that makes agreement between them evidence of anything. Some oracle binaries are IREE-only for that reason and refuse to start on the XLA backend. So a training run wants one lowerer and a full verification run wants both, and §B.4.2 is where the second one earns its place.

Selecting it is LEAN_MLIR_LOWERER=iree, and every tier of Getting started has an -iree twin that does so (lake run mnist-iree, and so on); the Imagenette twin runs five of the seven nets, since ResNet-50 and MobileNetV4 have no IREE driver. It additionally needs iree-compile on $PATH, the runtime built per historical/IREE_BUILD.md, and ffi/libiree_ffi.so. Data-parallel runs are XLA-only — collectives exist only on the PJRT path — so the ImageNet tier is not available on it.