B Getting started

If you want to just see it run, meaning train a real neural network end to end with the Lean \(\to \) MLIR \(\to \) GPU pipeline from Part 1, this appendix is the smallest path from zero to a trained model.

What you install depends on how far up the scale you want to go, and the steps are not evenly spaced:

Everything in Part 1 of this book is the middle row. It needs no Python at run time and no compiler: you download one shared library and run a Lean binary. Only the ImageNet tier (Track 4) needs an interpreter, and only to stream data.

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 PoC + SGD descent
lake build Proofs            # the whole suite
lake build Certs             # every certificate CI checks -- the long one

ProofsMinimal is the smallest thing that is still a real result: LinearFold (the rendered linear layer denotes the Mathlib fderiv math) and SgdDescentLinear (the step decreases the loss). Start there if you want to see one theorem land before committing to the full build.

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 C 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/ re-runs Lean’s kernel typechecker independently over 52 theorems, sandboxed, verifying each one’s axiom closure statically:

tests/comparator/run.sh

It splits in two. Challenge.lean imports Mathlib and nothing else. It carries the 13 architecture-free theorems (Chapter 1’s pdiv calculus rules and Chapter 9’s matrix-level ones), with the handful of definitions they need copied inline, and chk_pdiv_is_fderiv pinning pdiv to Mathlib’s fderiv by rfl. A reviewer can read that one file and know what is claimed without reading a line of this project. Copying the definitions is safe here in a way copying an architecture would not be: comparator compares the challenge and solution statements bit-identically, so a copy that drifted would fail the run rather than quietly prove something weaker.

The other 39 live in ChallengeArch.lean and do import the project, because each is a statement about a specific network, and “ResNet-34’s rendered backward equals its Fréchet derivative” cannot be phrased without ResNet-34 in scope. Once the fderiv pin and the structural rules are checked over Mathlib alone, the architecture theorems are applications of them, and what you must additionally trust is the forward functions and nothing else.

Track 2: Native install (XLA/PJRT)

This is the default path and the one every number in this book was measured on. There is no compiler to build.

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 -- the PJRT plugin ONLY
   pip install jax-rocm7-pjrt      # AMD    -- no jax, no jaxlib
   

   The shim looks under .venv/lib/python3.*/site-packages/jax_plugins/, which is where that command puts it, and $PJRT_PLUGIN points elsewhere. 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:

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

5. Build and run a trainer. No backend environment variables: XLA finds the device.

   lake build resnet34-verified-adam
   .lake/build/bin/resnet34-verified-adam data
   

   Use CUDA_VISIBLE_DEVICES / HIP_VISIBLE_DEVICES to pick a card, as with any GPU program.

Track 3: One-command demo tiers

With Track 2 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/Adam x bn/no-bn  (~19 min)
lake run imagenette   # the Part-I nets at 224^2, 80-epoch AdamW      (~7 h)

Those are XLA wall-clocks on one RTX 4060 Ti. The Imagenette figure is the five chapter nets; ResNet-50 (about 75 minutes) and MobileNetV4 (about 35) ride in the same tier and add to it. Each trainer’s transcript is teed to runs/<date>-<trainer>/. Each tier also has an -iree twin — see The second lowerer at the end. 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 and cifar8w-bn-ablation, each running SGD / momentum / AdamW in sequence — six arms in two 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

Track 4: ImageNet-1k runners

The scale tier. Six networks train on full ImageNet-1k off the same certified renderer as their Imagenette counterparts, in the chapter order §5 uses. A row is a (network, recipe) pair rather than a network: ResNet-50 appears twice because Chapter 5 trains it on two recipes, and the recipe is the middle field of the job’s own name.

Each job named here is the one that produced its chapter’s number. scripts/jobs/ also carries axis siblings that this table does not list — an fp32 peer of the ResNet-50 2018 row, an earlier A3 artifact — which differ from their row in precision or in one optimizer knob, never in the recipe. Each config’s header says which it is.

This tier needs Python

ImageNet data does not sit in a directory the way MNIST and Imagenette do. It is streamed from tfds through a generated per-net shim, a Python script emitted from that net’s own reference TrainConfig, so the augmentation the verified trainer consumes is provably the augmentation its JAX reference trains on. Generate them before the first run:

scripts/gen_shims.sh

Each net names its own shim, and the driver refuses to start if the named file is absent rather than falling back to another net’s, because an earlier fallback silently gave every net ResNet-34’s augmentation.

Running one

Single device, with the knobs the recipe needs:

LEAN_MLIR_VARIANT=mom256 LEAN_MLIR_BATCH=256 LEAN_MLIR_BASE_LR_U=100000 \
  .lake/build/bin/resnet34-imagenet-verified data

LEAN_MLIR_BATCH must match the batch its variant was rendered at: the batch is baked into the graph, not a runtime dimension, so a mismatch is a shape error at the first invoke. LEAN_MLIR_BASE_LR_U is the base learning rate in millionths (100000 = 0.1), integer-encoded because this toolchain has no String.toFloat?. The default 0.001 is an AdamW rate and will under-step a heavy-ball render by \(\sim \)100\(\times \), which looks like a broken render rather than a wrong knob. That is the one-card shape. The four-card one, with every knob taken from the job’s own config, is lake run r34-default-4gpu once — the next section.

Running one for a day and a half

These are multi-day runs, so they go through a supervisor (scripts/supervise.sh) that owns the device list, the epoch budget and the restart policy, and the job’s own name is the command:

lake run r34-default-4gpu        # supervised: resumes, restarts on AER / heat / stall
lake run r34-default-4gpu plan   # the precheck and the wall-clock on file; runs nothing
lake run r34-default-4gpu once   # one foreground run with the job's env, no restarts
lake run imagenet                # every row's plan, in chapter order; runs nothing
lake run imagenet start          # the rows in order, each resuming from its checkpoint

lake run imagenet is the tier itself, the fourth of the four. Bare, it only plans, because its twelve rows — the table above and the side quests below — are weeks of four-card time; start is the confirmation.

Data parallelism needs both replica knobs, and setting only one fails with a confusing “compiled for N replicas” error. PJRT_REPLICAS configures the shim’s client and compile options, and LEAN_MLIR_REPLICAS makes the Lean driver call the data-parallel invoke. The job configs set both. Collectives exist only on the PJRT path, so the data-parallel runs are XLA-only.

Side quests

Several chapters end in a variant that is rendered but is not a Track-4 runner, and the section heading says which of three kinds it is: a second recipe, a wider or deeper size, or an architecture the book reaches for once. They are collected here rather than in the table above because none of them has been trained — not because they cannot start. Five of the seven have job configs, launch like any Track-4 row and ride in lake run imagenet right after their chapter’s row; only the two ResNet-50 recipes have none. The status column says what each one is missing.

Common troubleshooting

• no lowerer shim could be loaded. No ffi/libpjrt_ffi.so. Build it with the gcc line in Track 2, which needs nothing but libc. (The same error names libiree_ffi.so when you have asked for that lowerer.)

• $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, as in Track 2, puts it where the shim looks by default.

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

• 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 §C.3.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 §C.3.2 is where the second one earns its place, including the asymmetry it leaves behind.

Selecting it is LEAN_MLIR_LOWERER=iree, and every Track 3 tier has an -iree twin that does so (lake run mnist-iree, and so on). 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 Track 4 is not available on it.