Installation

If you want to try TorchLean on a laptop, start with the CPU build. It does not require PyTorch, CUDA, or a GPU. The repository pins its Lean version in lean-toolchain, so Elan will select the right compiler for you.

A Five-Minute CPU Install

First install Elan, the Lean toolchain manager. On Linux or macOS:

curl https://elan.lean-lang.org/elan-init.sh -sSf | sh

Open a new terminal so that elan, lean, and lake are on your PATH. Then clone and build TorchLean:

git clone https://github.com/lean-dojo/TorchLean.git
cd TorchLean
lake exe cache get
lake build

The cache command downloads compatible prebuilt Lean dependencies when they are available. It is safe to omit; lake build will compile anything that is missing.

Run a small model to check the executable path:

lake exe torchlean quickstart_mlp --device cpu --steps 10

If those commands succeed, TorchLean is installed. You can inspect the available examples and verification commands with:

lake exe torchlean --help
lake exe verify --help

That CPU build is the common starting point on every platform. From there, TorchLean can compile its native CUDA runtime or link an external provider without changing the Lean model being run. The table below separates paths that work today from targets that are represented in the backend architecture but still need platform-specific runtime work.

Platform CPU NVIDIA GPU LibTorch provider Current status
Linux ✓ ✓ Native CUDA Standalone CUDA SDPA bridge test Supported; LibTorch eager dispatch is not wired
macOS, Intel or Apple silicon ✓ Not applicable Not yet CPU supported; Metal is planned
Windows with WSL2 ✓ Linux path ✓ CUDA on WSL2 Linux path Recommended Windows setup
Native Windows Bring-up target Not validated Not wired Backend target exists; native toolchain work remains

Here, “LibTorch provider” means the current scaled-dot-product-attention bridge, not a requirement for ordinary TorchLean models and not a claim that every operation is delegated to PyTorch. The CPU, native CUDA, and LibTorch sections below give the corresponding build commands.

Linux

CPU

You need Git, curl, and a C/C++ compiler. On Ubuntu or Debian:

sudo apt update
sudo apt install -y git curl build-essential

Then follow the five-minute install above. The default build uses the portable CPU runtime. It also builds harmless CUDA stub archives so that CPU-only machines can compile the complete Lean project; the stubs do not pretend that a GPU is present.

NVIDIA CUDA

Install a supported NVIDIA driver and CUDA toolkit using NVIDIA’s CUDA Installation Guide for Linux. TorchLean needs nvcc, cuBLAS, and cuFFT. Check the machine before rebuilding:

nvidia-smi
nvcc --version

Build and run the CUDA configuration:

lake -R -K cuda=true build
lake -R -K cuda=true exe torchlean quickstart_mlp \
  --device cuda --steps 10 --show-backend

The two CUDA choices happen at different times. -K cuda=true tells Lake to compile and link the native CUDA implementation. --device cuda asks the executable to use it. A CPU-linked executable rejects --device cuda instead of silently moving the run back to the CPU.

Use -R whenever you switch between CPU and CUDA configurations; it forces Lake to recompute the build description. The CUDA regression suite is:

lake -R -K cuda=true exe nn_tests_suite

The CUDA guide covers deterministic reductions, parity checks, sanitizers, and the remaining native-code trust boundary.

Optional LibTorch Attention

The normal CPU and CUDA builds do not need LibTorch. TorchLean currently uses LibTorch only through an optional scaled-dot-product-attention bridge. Download a matching GPU-enabled distribution from the official LibTorch installation page and extract it somewhere outside the repository.

The extracted directory must contain include/ and lib/. Pass its absolute path to Lake:

lake -R -K cuda=true -K libtorch=true \
  -K libtorch_home=/absolute/path/to/libtorch build
lake -R -K cuda=true -K libtorch=true \
  -K libtorch_home=/absolute/path/to/libtorch exe libtorch_sdpa_test

This enables one registered external provider. It does not route every TorchLean operation through PyTorch, and it does not make LibTorch part of Lean’s trusted kernel.

macOS

Install Apple’s command-line developer tools, then Elan and TorchLean:

xcode-select --install
curl https://elan.lean-lang.org/elan-init.sh -sSf | sh
git clone https://github.com/lean-dojo/TorchLean.git
cd TorchLean
lake exe cache get
lake build
lake exe torchlean quickstart_mlp --device cpu --steps 10

The CPU path works on Intel and Apple silicon. Modern macOS has no NVIDIA CUDA execution path. TorchLean already reserves --device metal in its device vocabulary, but Metal/MPS kernels are not implemented yet. Selecting Metal therefore returns an unsupported-device error; it does not quietly run the CPU implementation.

Windows

The most reliable Windows setup is Ubuntu under WSL2. Open an administrator PowerShell prompt:

wsl --install -d Ubuntu

After Windows restarts, open Ubuntu and follow the Linux instructions. For an NVIDIA GPU, follow NVIDIA’s CUDA on WSL guide. Install the Windows NVIDIA driver and the CUDA toolkit inside WSL; do not install a second Linux display driver inside WSL.

Native Windows

Native Windows is represented in the backend target vocabulary, but it is still a bring-up target rather than a regularly tested release path. Install Git, the Visual Studio C++ build tools, and the Windows SDK. Then install Elan from PowerShell:

curl -O --location https://elan.lean-lang.org/elan-init.ps1
powershell -ExecutionPolicy Bypass -f elan-init.ps1
del elan-init.ps1

The intended native CPU commands are:

lake exe cache get
lake build

The remaining work is platform engineering: the native libraries must be compiled with a compatible Windows C/C++ toolchain; CUDA and LibTorch must be discovered as .lib and DLL artifacts; Linux linker options such as -Wl,-rpath must be replaced; and the CPU stubs and GPU runtime must be tested under the Windows loader and ABI. Once those pieces are wired, the existing device, provider, and capsule abstractions do not need to be redesigned. Until then, WSL2 is the supported route for both CPU and NVIDIA GPU execution on Windows.

Use TorchLean From Another Lean Project

Add TorchLean to the downstream project’s lakefile.lean:

require TorchLean from git "https://github.com/lean-dojo/TorchLean.git" @ "main"

Then update and build:

lake update
lake exe cache get
lake build

Most model files need only:

import NN
open TorchLean

For development against a neighboring checkout, use a path dependency:

require TorchLean from "../TorchLean"

From A Model To A Kernel

Installation ends once the project builds, but accelerated execution introduces one more question: when a graph contains matmul, attention, or convolution, which implementation is allowed to run it?

TorchLean separates three names:

A device alone is not enough. Two CUDA providers may use different layouts, mask conventions, gradient implementations, or trust assumptions. TorchLean records those facts in a kernel capsule.

What A Kernel Capsule Looks Like

A capsule is a Lean record for one implementation of one operation. The native fused-attention capsule is defined in NN.Backend.Attention; the essential fields look like this:

def nativeFlashAttention : KernelCapsule :=
  { name := "native_cuda.flash_attention"
    op := .scaledDotProductAttention
    provider := .nativeCuda
    device := .cuda
    specName := "Spec.flashAttention"
    trustLevel := .checked
    supportsForward := true
    vjpMode := .backendVJP
    shapeContract := ...
    layoutContract := ...
    valueContract := ...
    vjpContract := ... }

Each contract field contains a structured claim rather than a free-form success label. For example, the layout field can state

{ claim := .layoutCompatibility .scaledDotProductAttention .flatRowMajor
  summary := "Q, K, V, mask, and output use contiguous row-major buffers"
  evidence := .runtimeGuard "Runtime.Autograd.Cuda.Tape.flashAttention" }

The claim says what must hold; the evidence says why the current profile is willing to use that implementation. A runtime guard or regression suite remains visible as a guard or test. It is not promoted to a theorem by placing it in the record.

The planner also checks that the shape, layout, value, and VJP fields contain the corresponding claim for the capsule’s operation. A proof attached to the wrong field is rejected. For theorem and checker evidence, the capsule author is still responsible for stating the exact Lean proposition represented by the structured claim; specName is an audit label, not a theorem by itself.

The record answers questions that are otherwise easy to lose in runtime plumbing:

Field Question it answers
op Which graph operation does this implementation claim to execute?
provider, device Which runtime family and hardware target may select it?
specName Which Lean-level meaning is it expected to refine?
shape and layout contracts Which dimensions, memory order, mask convention, and payload assumptions are required?
value contract What supports the forward-value claim?
VJP mode and contract Who computes the local gradient, and what supports that claim?
runtime support Is the path executable, test-only, planner-only, or not wired yet?
trust level Is the evidence proof-backed, checked, fuzzed, or an external assumption?

Capsules are declared in the source registry. They are not invented after a run to make its backend look acceptable. The planner selects an existing capsule whose device, provider, gradient mode, and trust level satisfy the requested profile. Missing or inadmissible capsules cause planning to fail.

Where The Theorems Enter

A capsule does not prove foreign code merely by containing the name of a theorem. The proof itself must be a Lean declaration whose proof term is checked by Lean’s kernel. The capsule records which evidence is meant to support each obligation, and the audit and gate layers make missing or external evidence visible.

FlashAttention gives a concrete example. TorchLean already proves a semantic equality between its fused attention specification and standard scaled dot-product attention:

#check Spec.flashAttention_eq_scaledDotProductAttention
#check Spec.flashAttentionBackward_eq_scaledDotProductAttentionBackward

Those are genuine Lean theorems about two Lean definitions. They justify replacing the standard attention specification with the fused specification. They do not prove that a particular CUDA binary implements either definition. The current native CUDA capsule therefore records runtime guards, regression tests, source provenance, and trust level checked, not verified.

To move a native capsule to a stronger status, the work has to happen in this order:

  1. define the mathematical operation and its domain assumptions in Lean;
  2. state the runtime refinement claim, including shape, layout, Float32, and error assumptions;
  3. prove that claim, or build a replay checker with a soundness theorem;
  4. attach that theorem or checker to the capsule;
  5. keep any remaining FFI, compiler, driver, or hardware assumption explicit.

Proof and checker evidence carry their Lean proof terms. Source files, native symbols, runtime guards, regression suites, and fuzz runs are recorded separately. They remain useful audit information, but strict acceptance does not treat them as theorems.

Planning And Trust Policies

For each backend-visible graph operation, TorchLean follows the same sequence:

  1. Read the graph. Keep the operation, shape, payload, and proof-facing specification.
  2. Apply the profile. Restrict the device, providers, gradient ownership, and accepted trust levels.
  3. Select a capsule. Choose a compatible implementation already present in the registry.
  4. Audit and gate. Reject missing evidence or forbidden external boundaries.
  5. Admit, execute, and report. The eager runtime rejects non-eager capsules, consumes an accepted executable capsule, and prints it on first use when `--show-backend` is set.

An accepted kernel or graph plan carries a Lean proof that its selected policy gate returned accepted. No-grad execution requests only a forward contract; training rejects a forward-only capsule for differentiable operations. Random sources are non-differentiable and therefore do not need a local VJP.

The maintained profiles currently have these meanings:

Profile Forward provider Backward ownership Current status
checked_cpu Reference CPU runtime TorchLean tape Implemented
checked_cuda Native CUDA, cuBLAS, cuFFT TorchLean tape or named backend VJP Implemented
libtorch_forward_cuda Selected LibTorch forward capsules TorchLean tape Registered selectively; not a universal dispatcher
libtorch_autograd_cuda Selected LibTorch capsules External autograd Explicit larger trust boundary
future_* Metal, ROCm, WebGPU, XLA, Neuron, custom providers Provider-specific Extension targets, not implemented runtimes

The default CUDA profile does not silently choose LibTorch. LibTorch SDPA remains testOnly until its capsule is connected to eager multi-head attention. External providers must be enabled explicitly and remain visible in the backend report.

The corresponding Lean definitions are in NN.Backend.Capsule, NN.Backend.Profile, and NN.Backend.Gate.

Adding Another Platform

Metal, ROCm, WebGPU/WASM, TPU/XLA, Trainium/Neuron, custom chips, and caller-supplied accelerators all use the same extension points:

  1. add the device and provider vocabulary;
  2. detect build and runtime availability honestly;
  3. register capsules only for the operations the provider implements;
  4. connect the selected capsules to executable runtime dispatch;
  5. supply shape, layout, value, and VJP evidence at the right trust level;
  6. add platform CI and parity tests before calling the target supported.

Until those pieces exist, target selection fails. Asking for Metal or TPU must never become an unreported CPU run.

Check An Installation

These commands cover the normal CPU installation:

lake build
lake lint
lake exe nn_tests_suite
lake exe torchlean --help
lake exe verify --help

For CUDA, rebuild and run the suite with -R -K cuda=true.

For a complete account of Lean axioms, executable checkers, CUDA and FFI code, external artifact producers, and floating-point assumptions, read TRUST_BOUNDARIES.md.

References