TorchLean

3.1. Choosing How A Model Runs🔗

The MLP trained in A Tour Of The API maps a two-component input to a one-component output, with parameters \theta:

F_\theta:[2]\to[1]

Its type stays [2] → [1] whether it runs on the CPU, with FloatLib's bit-level binary32, or with LibTorch's CUDA operations. Every run records the same autograd tape; the device only chooses who executes each operation. Holding the architecture, seed, and data fixed lets us compare these choices on the same training problem. The transcripts below record the binary's output on the machine used for this guide; the embedded Lean comparisons also run when the page is built.

Runtime Configuration🔗

TorchLean configures arithmetic, providers, and model mode separately. The runtime validates their compatibility; not every combination is supported. Each setting controls a different part of the computation:

Dial

Values

What it decides

arithmetic

native, ieee

which binary32 semantics every scalar operation obeys

device and providers

cpu, gpu, a full profile

which implementation of each operation is selected, and what evidence backs it

train or evaluation mode

nn.Mode.train, nn.Mode.eval

how mode-sensitive layers such as dropout behave

Changing the device selects implementations of the model's operations. Changing train or evaluation mode can change the operations' behavior, as dropout will show below. Recording both choices is necessary to distinguish a provider comparison from a comparison of different forward computations.

-- The dispatcher vocabulary is broader than the trainer's -- accepted arithmetic. inductive TorchLean.Runtime.Arithmetic : Type number of parameters: 0 constructors: TorchLean.Runtime.Arithmetic.native : Runtime.Arithmetic TorchLean.Runtime.Arithmetic.ieee : Runtime.Arithmetic TorchLean.Runtime.Arithmetic.complex : Runtime.Arithmetic#print Runtime.Arithmetic
inductive TorchLean.Runtime.Arithmetic : Type
number of parameters: 0
constructors:
TorchLean.Runtime.Arithmetic.native : Runtime.Arithmetic
TorchLean.Runtime.Arithmetic.ieee : Runtime.Arithmetic
TorchLean.Runtime.Arithmetic.complex : Runtime.Arithmetic

The supervised trainer accepts two of these three constructors. As the errors below show, the lower dispatcher can execute complex binary32, but the supervised interface exchanges real data and results. Complex training instead uses an explicit real loss with autograd.complex.grad, followed by nn.sgdStep on complex state. The complex_regression command demonstrates this separate path; it preserves both components in predictions and checkpoints.

The device vocabulary is the widest of the three:

-- A device constructor names a target without asserting -- that a handler is linked. inductive NN.Backend.Device : Type number of parameters: 0 constructors: NN.Backend.Device.cpu : NN.Backend.Device NN.Backend.Device.cuda : NN.Backend.Device NN.Backend.Device.rocm : NN.Backend.Device NN.Backend.Device.metal : NN.Backend.Device NN.Backend.Device.wasm : NN.Backend.Device NN.Backend.Device.tpu : NN.Backend.Device NN.Backend.Device.trainium : NN.Backend.Device NN.Backend.Device.custom : NN.Backend.Device NN.Backend.Device.external : NN.Backend.Device#print Runtime.Device
inductive NN.Backend.Device : Type
number of parameters: 0
constructors:
NN.Backend.Device.cpu : NN.Backend.Device
NN.Backend.Device.cuda : NN.Backend.Device
NN.Backend.Device.rocm : NN.Backend.Device
NN.Backend.Device.metal : NN.Backend.Device
NN.Backend.Device.wasm : NN.Backend.Device
NN.Backend.Device.tpu : NN.Backend.Device
NN.Backend.Device.trainium : NN.Backend.Device
NN.Backend.Device.custom : NN.Backend.Device
NN.Backend.Device.external : NN.Backend.Device

Two of these nine devices currently have maintained runtime profiles. Device names also index profiles, capsules, and reports, so the vocabulary includes targets such as rocm and metal without claiming that an executable provider is available for them.

withDevice checks for a maintained profile before execution. validateForExecution additionally checks whether the linked binary can provide the requested runtime:

-- Pure profile lookup and effectful execution validation -- answer different questions. open Runtime.Autograd.Torch in Config.withDevice : Runtime.Autograd.Torch.Config → NN.Backend.Device → Except String Runtime.Autograd.Torch.Config#check @Config.withDevice open Runtime.Autograd.Torch in Config.validateForExecution : Runtime.Autograd.Torch.Config → IO Unit#check @Config.validateForExecution
Config.withDevice : Runtime.Autograd.Torch.Config → NN.Backend.Device → Except String Runtime.Autograd.Torch.Config
Config.validateForExecution : Runtime.Autograd.Torch.Config → IO Unit

The first returns Except String Config, so the failure is a value and can be inspected without running anything. Checking the four named devices takes four lines:

-- Check maintained defaults without requiring the named -- devices to be present. open Runtime.Autograd.Torch in cpu: maintained profile cuda: maintained profile metal: device `metal` has no maintained runtime profile; provide a backend profile with executable capsules tpu: device `tpu` has no maintained runtime profile; provide a backend profile with executable capsules #eval do let devices : List Runtime.Device := [.cpu, .cuda, .metal, .tpu] devices.forM fun d => do let name := d.cliName match Config.withDevice {} d with | .ok _ => IO.println s!"{name}: maintained profile" | .error msg => IO.println s!"{name}: {msg}"
cpu: maintained profile
cuda: maintained profile
metal: device `metal` has no maintained runtime profile; provide a backend profile with executable capsules
tpu: device `tpu` has no maintained runtime profile; provide a backend profile with executable capsules

A custom backend needs both a profile describing admissible kernels and runtime handlers that execute them. Supplying capsule metadata alone does not implement Metal. The maintained profiles reject unsupported devices at configuration time.

The Except String Config result lets a caller handle a missing profile before entering IO. Execution validation then checks the linked runtime and device availability. Individual operations still have their own shape, layout, arithmetic, and gradient requirements.

Runtime Flags🔗

The runner lists its examples and runtime flags:

Terminal
# List the runnable examples and the shared entry-point
# options.
scripts/lake.sh exe torchlean --help

The help lists the runnable examples, followed by their shared runtime flags:

Runtime flags:
  --choose                         ask for a device when the example supports --device
  --device auto|cpu|gpu|cuda|rocm|metal|wasm|tpu|trainium|custom|external
      execution backends: CPU and CUDA; other parsed device names are not implemented
  --arithmetic native|ieee|complex
      arithmetic availability depends on the example; check its --help
  --seed N
  --show-backend

Verification commands live under `scripts/lake.sh exe verify -- list`.
Use `scripts/lake.sh exe torchlean <example> --help` for command-specific flags;
`jaxpr` has no --help and takes only artifact paths.

The top-level parser knows three arithmetics; each example states which ones it supports. For our example:

Terminal
# Inspect the controls accepted by this particular training
# example.
scripts/lake.sh exe torchlean quickstart_mlp --help
TorchLean simple MLP quickstart

Usage:
  scripts/lake.sh exe torchlean quickstart_mlp [--steps N] [--seed S]
    [--arithmetic native|ieee] [--device cpu|cuda]
    [--show-backend]

This example supports native|ieee and cpu|cuda. An opt-in interactive chooser, --choose, can also select the device; Command-Line Reference shows its prompt.

The chooser and flags both request a configuration. The banner echoes that request; the backend report below names the implementations selected for the model's operations.

CPU Execution🔗

The baseline uses native binary32 arithmetic on the CPU and twenty training steps:

Terminal
# Run the short CPU baseline and print the selected backend
# capsules.
scripts/lake.sh exe torchlean quickstart_mlp \
  --device cpu \
  --steps 20 \
  --seed 2026 \
  --show-backend

Leaving out --show-backend for a moment, the run prints:

Recorded output
== Quickstart: simple MLP training (seed=2026, steps=20) ==
target(heldout)    = [0.200000]
untrained(heldout) = [-0.088261]
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.401184
steps=20 arithmetic=native scalar=Float32 loss=0.495227 -> 0.401184
trained(heldout) = [0.365380]

Training opens one session and records each sample's operations on a fresh tape as the model runs. Every operation asks the active profile for an admissible capsule, executes that capsule's provider, and appends a local VJP rule when gradients are required. A capsule bundles a provider, a trust level, a VJP owner, and a numerical policy, and --show-backend prints one the first time each operation is selected. Here is the entry for addition, wrapped to fit the page (in the terminal each of these is a single line):

  add: reference.add provider=reference trust=checked
    vjp=torchlean-tape reduction=n/a
    shape: shape safety for add; guarded at runtime by
      portable runtime shape checks
    layout: canonical-tensor layout compatibility for add;
      guarded at runtime by typed tensor layout
    value: add forward refines its TorchLean semantics;
      covered by test suite NN.Tests.Runtime.Floats.Suite
    vjp: add torchlean-tape VJP refines its TorchLean
      semantics; covered by test suite
      NN.Tests.Runtime.Floats.Suite

The four evidence lines distinguish shape, layout, forward-value, and VJP claims. Two cite runtime guards and two cite a test suite. These labels identify the declared evidence for the selected capsule; they do not assert a theorem. Verification And Certificates develops the proof interfaces.

Our twenty-step CPU run selects seven capsules. Tabulating the first line of each entry:

Operation

Capsule

VJP owner

Reduction policy

reshape

reference.reshape

torchlean-tape

n/a

permute

reference.permute

torchlean-tape

n/a

matmul

reference.matmul

torchlean-tape

fixed-left

broadcast

reference.broadcast

torchlean-tape

n/a

add

reference.add

torchlean-tape

n/a

relu

reference.relu

torchlean-tape

n/a

mse_loss

reference.mse_loss

torchlean-tape

fixed-left

Every provider is reference, every trust level is checked, and every VJP is owned by TorchLean's own tape. The fixed-left policy on the two reducing operations specifies their summation order, one of the conditions needed to reproduce floating-point results. The CUDA profile below changes that policy.

With --show-backend, the run prints the banner [TorchLean] backend capsules used: several times. An eager session prints it when it is created (Session/Lifecycle.lean), so the banners count sessions rather than operations. trainer.predict and trainer.train each open a trainer session. A trainer session creates one eager session that records a fresh tape for each training sample, and further eager sessions without gradients for prediction and evaluation loss. Their banners print together when the trainer session opens, so a banner can have nothing under it. The banner therefore identifies a session opening; the entries beneath it identify capsule selections.

The tape lets you inspect provider choices and use the maintained CUDA runtime. A manually assembled eager session can record a different operation sequence on each call. The current model frontend still hides tensor elements behind Function, so the tape alone does not enable arbitrary tensor-dependent Lean branches. Dynamic Control Flow And Typed Graphs below explains this frontend boundary.

There are three different loss observations in this short run. The initial 0.495227 is an average over the 25 training examples. The step 0 value, 0.250488, belongs to the example used for that update. The final 0.401184 is another evaluation over the dataset after 20 updates. Comparing the first and last values is meaningful as a change in dataset loss; comparing either one directly with the single-example step loss mixes two different measurements. Twenty updates also need not visit every example equally often.

The held-out prediction is a fourth observation. Its target is 0.2, and its trained value here is 0.365380; a lower training loss does not make that particular prediction exact. Keeping it separate makes the log useful even when individual update losses fluctuate.

The seven capsule kinds describe the operations needed by the model and loss. They are not a count of kernel launches or arithmetic instructions. Repeated operations may choose the same capsule, and a report with no new entries during prediction does not mean prediction performed no work. It means that report contains no additional selections to display for that session.

Training Inside The Guide🔗

Comparisons can also run inside the guide, so the page build checks them directly. The setup fixes the quickstart model, target, dataset, optimizer, and seed:

-- Hold model, target, data, and seed fixed while exposing -- the runtime choices. /-- The `2 -> 8 -> 1` regression MLP from the quickstart. -/ def runModel : nn.Builder (nn.Sequential [2] [1]) := nn.Sequential![nn.linear 2 8, nn.relu, nn.linear 8 1] /-- Target `0.8 relu(x₁+x₂) - 0.4 relu(x₂-x₁) + 0.2`. -/ def runTarget (x : Tensor Float [2]) : Tensor Float [1] := let relu (v : Float) := if v < 0.0 then 0.0 else v let rise := 0.8 * relu (x[0] + x[1]) let fall := 0.4 * relu (x[1] - x[0]) [rise - fall + 0.2] def runInputs : Tensor Float [25, 2] := Data.Synthetic.squareGrid (-1.0) 1.0 5 def runData : Trainer.Dataset [2] [1] := Data.fromTensors runInputs (Tensor.mapLeading [25] runTarget runInputs) /-- A grid point that is not in the training set. -/ def runPoint : Tensor Float [2] := [0.25, -0.75] /-- One trainer per choice of arithmetic and device. -/ def trainerOn (arithmetic : Runtime.Arithmetic) (device : Runtime.Device) : Trainer [2] [1] := Trainer.new runModel { objective := mse optimizer := optim.adam { lr := 0.03 } arithmetic := arithmetic device := device seed := 2026 }

trainerOn exposes arithmetic and device as arguments while holding the training problem fixed. Each call opens a fresh run from seed 2026, so two runs differ only in the argument that changed, not in parameters left over from an earlier run.

CUDA Execution Through LibTorch🔗

CUDA support is a build-time choice. Set TORCHLEAN_LIBTORCH_HOME to a CUDA-enabled SDK as described in the installation guide, then select it with -K cuda=true:

Terminal
# Select the CUDA build profile before requesting a CUDA
# training session.
scripts/lake.sh -K cuda=true exe torchlean quickstart_mlp \
  --device cuda \
  --steps 20 \
  --seed 2026 \
  --show-backend

The report now names a different provider for every operation. Comparing it against the CPU table:

Operation

CPU capsule

CUDA capsule

CUDA reduction policy

reshape

reference.reshape

libtorch.reshape

n/a

permute

reference.permute

libtorch.permute

n/a

matmul

reference.matmul

libtorch.matmul

implementation-defined

broadcast

reference.broadcast

libtorch.broadcast

n/a

add

reference.add

libtorch.add

n/a

relu

reference.relu

libtorch.relu

n/a

mse_loss

reference.mse_loss

libtorch.mse_loss

implementation-defined

The provider label changes from reference to libtorch. Evidence names the retained CUDA comparison for each tested operation; operations without one record a LibTorch trust boundary. Shape guards move to CUDA FFI size and rank checks at the Lean and native boundary. The default GPU profile admits those external implementation assumptions; the strict checked profile does not.

Local VJP ownership depends on the operation. ReLU and broadcast use backend-vjp: their derivatives call native backward primitives. Matmul, addition, MSE, reshape, and permutation retain torchlean-tape: Lean composes their derivatives from LibTorch numerical operations. TorchLean records and traverses the tape in both cases; LibTorch autograd recording is disabled. For numerical comparison, the two reducing operations move from fixed-left to implementation-defined.

To compare the two devices, hold the model, dataset, seed, optimizer, and step count fixed. A provider with reduction=implementation-defined may use a different summation order, and floating point addition is not associative (Goldberg, 1991). Compare the artifact relevant to the experiment, such as predictions or parameter buffers, and record its dtype, provider, and comparison tolerance. Agreement in printed decimal output does not establish equal parameter bits or equal behavior on other inputs.

The module and CUDA buffer both use binary32 values, but upload is not conversion-free. The current bridge widens host Float32 elements to Float staging values before packing them into a CUDA float32 buffer. This preserves finite binary32 values but adds transfer work. A capsule identifies the trusted native provider; it does not prove the kernel, compiler, driver, or device correct.

A build without CUDA support rejects the request. The default CPU profile supports the portable runtime without a GPU; compiling it does not exercise the CUDA provider.

A common seed fixes initialization and the example sequence; it does not fix the order in which a backend adds partial sums. Updates can carry numerical differences into later steps, and a ReLU can change its active branch near zero. To investigate a discrepancy, compare the inputs and outputs of the first differing operation before attributing the final loss difference to the optimizer.

Executable Binary32 Arithmetic🔗

Terminal
# Use executable binary32 arithmetic for the same short CPU
# workload.
scripts/lake.sh exe torchlean quickstart_mlp \
  --arithmetic ieee \
  --device cpu \
  --steps 2 \
  --seed 2026

The ieee path evaluates addition and multiplication with FloatLib binary32. Comparing the two implementations can expose errors in rounding or exceptional-value handling that a comparison of decimal output alone may miss. Formal floating-point models make these rules explicit (Goldberg, 1991). Flocq made such a model reusable inside a proof assistant (Boldo and Melquiond, 2011), and the same group's later work connects such a model to the arithmetic a compiler actually emits (Boldo et al., 2015). FloatLib supplies the executable format and arithmetic used here. For higher precision, select a valid FloatLib binary format in typed CPU tensors and graphs; the trainer flag shown above remains fixed to binary32.

For typed training, trainer.openTyped (α := α) keeps inputs, state, predictions, and losses in the chosen scalar type on CPU. It rejects GPU devices and custom backend profiles. Supply initialState? when the parameters need precision beyond the seeded Float initializers; converting those initializers into a wider format does not add precision. Optimizer and scheduler coefficients are still configured in Float and converted into the selected scalar.

The following embedded comparison runs two hundred steps with each arithmetic choice, holding the seed and dataset fixed. It also checks prediction equality:

-- Keep the device fixed to compare the two binary32 -- implementations. dataset size = 25 mean_loss(before training) = 0.495227 mean_loss(after training) = 0.002402 dataset size = 25 mean_loss(before training) = 0.495227 mean_loss(after training) = 0.002402 native = steps=200 arithmetic=native scalar=Float32 loss=0.495227 -> 0.002402 ieee = steps=200 arithmetic=ieee scalar=ExecFloat.Binary 8 23 loss=0.495227 -> 0.002402 native prediction = [0.228325] ieee prediction = [0.228325] predictions equal = true #eval do let opts : Trainer.TrainOptions := { steps := 200, logEvery := 0 } let native ← (trainerOn native cpu).train runData opts let ieee ← (trainerOn ieee cpu).train runData opts IO.println s!"native = {native.summary}" IO.println s!"ieee = {ieee.summary}" let a ← native.predict runPoint let b ← ieee.predict runPoint IO.println s!"native prediction = {a}" IO.println s!"ieee prediction = {b}" IO.println s!"predictions equal = {a == b}"
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.002402
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.002402
native = steps=200 arithmetic=native scalar=Float32 loss=0.495227 -> 0.002402
ieee   = steps=200 arithmetic=ieee scalar=ExecFloat.Binary 8 23 loss=0.495227 -> 0.002402
native prediction = [0.228325]
ieee   prediction = [0.228325]
predictions equal = true

The held-out target can be checked without running either trainer. At [0.25, -0.75], the sum of the coordinates is -0.5 and their difference x₁ - x₀ is -1. Both ReLU terms in the target function therefore vanish, leaving its constant 0.2. The common prediction 0.228325 is about 0.028325 above that target. The small final dataset loss and this remaining prediction error answer different questions, even though both came from the same trained parameters.

Agreement on one model, one seed, and one dataset is a test result. A differential test between a hardware instruction and a bit-level implementation can expose rounding, tie, or sign errors exercised by its inputs and assertions. It says nothing about inputs we did not try. The theorem-side story lives in Floating Point Semantics, where the proof-oriented FP32 model and exact Real appear in statements rather than in IO, and where the finite and no-overflow bridge from FP32 to FloatLib binary32 is stated. FP32 rounds reals at binary32 precision with gradual underflow, and it does not model overflow, NaN, infinity, or signed zero, which is why the bridge has hypotheses.

The bit-level implementation can cost more than native arithmetic. Measure that cost on the workload and hardware of interest, separating process startup from repeated execution. Record commands, software versions, thread settings, and the actual provider alongside timings so that others can reproduce the comparison.

Keep the training horizon fixed when comparing arithmetic choices. A two-step demonstration and a two-hundred-step run answer different questions about the same model.

There is also an interface boundary in these examples: the trainer accepts and reports public Float values while its selected implementation works with a binary32 scalar representation. An explicit Float autograd program elsewhere in the guide is a separate arithmetic choice. Matching the public input type alone does not establish matching internal arithmetic.

Runtime Configuration In Lean🔗

Trainer.RunConfig exposes the runtime choices as typed fields:

-- Record updates isolate runtime changes while retaining -- the optimizer settings. def runCpu : Trainer.RunConfig := { arithmetic := native device := cpu optimizer := optim.adam { lr := 0.03 } } def runIeeeCpu : Trainer.RunConfig := { runCpu with arithmetic := ieee } def runCuda : Trainer.RunConfig := { runCpu with device := gpu } /-- Bind a run configuration to an objective and a seed. -/ def trainerFromRun (run : Trainer.RunConfig) : Trainer [2] [1] := Trainer.new runModel { run with objective := mse, seed := 2026 } cpu native cpu ieee cuda native #eval do let runs := [runCpu, runIeeeCpu, runCuda] runs.forM fun (run : Trainer.RunConfig) => do IO.println s!"{run.toConfig.deviceName} {run.arithmetic.cliName}"
cpu native
cpu ieee
cuda native

The record update syntax makes runIeeeCpu and runCuda each differ from runCpu in exactly one field. RunConfig extends Runtime.Config with the optimizer and arithmetic, and Trainer.Config extends RunConfig with the objective, so one record update such as { run with objective := mse, seed := 2026 } produces a full Trainer.Config. This permits reuse of one runtime configuration across several objectives. run.toConfig is the Runtime.Config that the session layer actually validates, which is the value that validateForExecution inspects.

The same syntax goes the other way: { runtime with optimizer := ... } starts a RunConfig from a Runtime.Config parsed from the command line. A sweep therefore declares the model once, keeps one base RunConfig, and derives each cell of the sweep from it.

Constructing runCuda does not open CUDA. These labels print only the device and arithmetic fields of each record; they omit optimizer settings and provider overrides. Save those fields too when comparing runs.

Graph Lowering And Direct Execution🔗

The trainer runs on the tape and keeps no graph. When a program needs a typed graph, for example to take forward-mode derivatives, lower the model directly:

-- Graph leaves contain state followed by input; the wrapper -- separates them for callers. @Runtime.Autograd.Model.Layers.Seq.lowerToTypedGraph : {σ τ : Shape} → (model : Runtime.Autograd.Model.Layers.Seq σ τ) → optParam Runtime.Autograd.Model.Layers.Mode Runtime.Autograd.Model.Layers.Mode.eval → {α : Type} → [inst : Storage α] → [Context α] → IO (Runtime.Autograd.Torch.TypedGraph α (model.stateShapes ++ [σ]) τ)#check @nn.lowerToTypedGraph @nn.TypedGraphModel.forward : {σ τ : Shape} → {α : Type} → [inst : Storage α] → {stateShapes : List Shape} → nn.TypedGraphModel stateShapes σ τ α → nn.State α stateShapes → Tensor α σ → Tensor α τ#check @nn.TypedGraphModel.forward
@Runtime.Autograd.Model.Layers.Seq.lowerToTypedGraph : {σ τ : Shape} →
  (model : Runtime.Autograd.Model.Layers.Seq σ τ) →
    optParam Runtime.Autograd.Model.Layers.Mode Runtime.Autograd.Model.Layers.Mode.eval →
      {α : Type} →
        [inst : Storage α] → [Context α] → IO (Runtime.Autograd.Torch.TypedGraph α (model.stateShapes ++ [σ]) τ)
@nn.TypedGraphModel.forward : {σ τ : Shape} →
  {α : Type} →
    [inst : Storage α] →
      {stateShapes : List Shape} →
        nn.TypedGraphModel stateShapes σ τ α → nn.State α stateShapes → Tensor α σ → Tensor α τ

lowerToTypedGraph returns a TypedGraph α (model.stateShapes ++ [σ]) τ, whose index is a flat list of every graph leaf: parameters first, then the model input. nn.TypedGraphModel is the model-facing view of exactly that type, and its forward splits the leaves back into nn.State α stateShapes and Tensor α σ. Nothing is converted; the second signature is a way of talking about the first that keeps the parameters and the input apart.

nn.TypedGraphModel is an abbrev, so it is transparent to dot notation. Writing g.forward state input resolves to the raw TypedGraph.forward, which wants one TensorPack over the concatenated leaves and will reject the two arguments. Call the model-facing wrappers by name, as below.

-- Zero parameter tangents isolate derivatives with respect -- to the input coordinates. forward = [-0.088261] input gradient = [0.263189, 0.205411] parameter shapes = [[8, 2], [8], [1, 8], [1]] jvp along e₀ = [0.263189], jvp along e₁ = [0.205411] #eval do let built := nn.build 2026 runModel let g ← nn.lowerToTypedGraph built (α := Float) let state := nn.initialState built let shapes := nn.stateShapes built let y := nn.TypedGraphModel.forward g state runPoint IO.println s!"forward = {y}" let (_, inputGradient) := nn.TypedGraphModel.vjp g state runPoint [1.0] IO.println s!"input gradient = {inputGradient}" IO.println s!"parameter shapes = {shapes}" let zero : nn.State Float shapes := nn.State.zeros let jvp := nn.TypedGraphModel.jvp g state zero runPoint let d0 := jvp [1.0, 0.0] let d1 := jvp [0.0, 1.0] IO.println s!"jvp along e₀ = {d0}, jvp along e₁ = {d1}"
forward = [-0.088261]
input gradient = [0.263189, 0.205411]
parameter shapes = [[8, 2], [8], [1, 8], [1]]
jvp along e₀ = [0.263189], jvp along e₁ = [0.205411]

forward = [-0.088261] is the same number the runner printed as untrained(heldout) in every experiment above, which is a small consistency check that the lowered graph and the trainer agree on the model. The parameter shapes [[8, 2], [8], [1, 8], [1]] are the two weight matrices and two bias vectors of a 2 → 8 → 1 MLP, in declaration order.

The reverse mode pass, seeded with the output cotangent [1.0], returns the input gradient [0.263189, 0.205411]. The two forward mode passes, seeded with the basis tangents [1.0, 0.0] and [0.0, 1.0] and a zero parameter tangent, return [0.263189] and [0.205411]. For a function with a one-dimensional output these must coincide, because both are computing the same 1 \times 2 Jacobian, one row at a time from the left and one column at a time from the right. They coincide to the last printed digit. Comparing forward and reverse directional derivatives checks their agreement without requiring a separate symbolic derivative (Baydin et al., 2018). Retaining the graph lets all three calls use the same recorded operations.

Two lifetimes to keep straight. nn.lowerToTypedGraph hands you a graph that lives as long as the reference does. An eager session instead records one tape per recording phase, and resetTape starts a fresh one. When the graph itself must survive across calls, lower it.

The example explicitly uses Float; it should not be read as another execution of the trainer's binary32 comparison, despite the similar printed initial prediction.

Device And Provider Profiles🔗

Ordinary code selects a device and lets TorchLean resolve the maintained profile for it. Advanced code can select a complete profile with withBackendProfile, choosing provider preference, assurance policy, VJP ownership, target availability, and capsule modules together. Calling withDevice afterwards clears that override, which is deliberate: a device name and a hand-built profile are two ways to answer the same question, and the last answer wins rather than silently merging.

Backend Selection follows an operation request through provider selection, contract checks, and binding to a runtime handler.

Train And Evaluation Mode🔗

The high-level Trainer.Session selects mode from the operation you call. Lower-level module APIs also provide explicit mode setters:

Call

Mode

session.step sample

training

session.step batch (batch := true)

training

session.predict input

evaluation

session.loss sample

evaluation

session.eval data

evaluation

trainer.train uses the training path for updates and the evaluation path for the losses it reports, for the summary prediction, and for later calls to Trainer.Result.predict. Mode is independent of device and execution choice: a CUDA session switches mode without changing the model architecture or the provider profile.

step returns Unit by default, so an update does not read the loss back to the host. Add (loss := true) when the caller needs the loss from that update. The batch flag selects a nonempty array of samples; it does not change training mode.

Dropout (Srivastava et al., 2014) makes the effect of mode visible. During training it samples a mask; during evaluation it applies its deterministic inference behavior. The next block inserts dropout between the MLP's hidden activation and output layer, then evaluates the same dataset through both session paths:

-- Evaluate before the batch update so the mode comparison -- starts at the same state. /-- The same MLP with one dropout layer. -/ def runDropoutModel (p : Float) : nn.Builder (nn.Sequential [2] [1]) := nn.Sequential![nn.linear 2 8, nn.relu, nn.dropout p, nn.linear 8 1] def runDropoutTrainer (p : Float) : Trainer [2] [1] := Trainer.new (runDropoutModel p) { objective := mse optimizer := optim.adam { lr := 0.03 } seed := 2026 } p = 0.000000 evaluation mode 0.495227 training mode 0.495227 gap * 1e9 = -29.355288 p = 0.500000 evaluation mode 0.761530 training mode 0.307745 gap * 1e9 = -453785276.710987 #eval do let batch := (← runData.materialize (α := Float)).toArray for p in [0.0, 0.5] do let session ← (runDropoutTrainer p).open let evaluation ← session.eval runData let training ← session.step batch (batch := true) (loss := true) IO.println s!"p = {p}" IO.println s!" evaluation mode {evaluation}" IO.println s!" training mode {training}" let gap := (training - evaluation) * 1e9 IO.println s!" gap * 1e9 = {gap}"
p = 0.000000
  evaluation mode 0.495227
  training mode   0.495227
  gap * 1e9 = -29.355288
p = 0.500000
  evaluation mode 0.761530
  training mode   0.307745
  gap * 1e9 = -453785276.710987

Read the gaps, not the printed losses. At p = 0.0 the layer is the identity in both modes and the two numbers print the same to six decimals, yet the gap is -2.9e-8 rather than zero. eval converts per-sample losses to Float and averages them there; step with (batch := true) uses the runtime scalar for its batch reduction. This comparison changes both reduction structure and the precision of the final average, so it does not isolate summation order alone.

At p = 0.5 the much larger gap illustrates the effect of training-mode dropout in this example. The p = 0 result is a useful control, not a universal bound on floating point discrepancies.

Repeating the p = 0.5 measurement in freshly opened sessions gives the same result for this seed. The builder derives the dropout layer's seed from the model build seed. This demonstrates seeded reproducibility; it does not establish independence of masks across calls. Session.save saves model state; it does not save the optimizer history, loader position, or random stream needed for exact training resume. The lower-level Checkpoint.Optimizer.save separately supports backend-owned CUDA Adam/AdamW state and its eager random counter.

Compare PyTorch, where the same distinction exists but is carried by a mutable attribute (Paszke et al., 2019):

>>> # Compare repeated calls before and after disabling
>>> # training-mode dropout.
>>> torch.manual_seed(2026)
>>> x = torch.tensor([0.25, -0.75])
>>> net = torch.nn.Sequential(
...     torch.nn.Linear(2, 8), torch.nn.ReLU(),
...     torch.nn.Dropout(0.5), torch.nn.Linear(8, 1))
>>> net.training
True
>>> [net(x).item() for _ in range(3)]
[-0.8670920133590698, 0.04916824400424957, -0.2292545586824417]
>>> net.eval()
>>> [net(x).item() for _ in range(3)]
[-0.1853325217962265, -0.1853325217962265, -0.1853325217962265]

Three calls to the same object on the same input give three different answers, then the same call becomes deterministic after net.eval(). In PyTorch, mode is a flag on the module and defaults to training. In TorchLean, Trainer.Session chooses the mode for these operations, while lower-level APIs require the caller to select it. The example derives training-mode randomness from the build seed.

The mode also determines which random state the reverse rule must retain. Consider the operation

y=\operatorname{Dropout}_{p}(x).

During training a mask is realized and retained for the backward rule. Re-running the backward pass with a freshly sampled mask would not differentiate the forward value that was actually computed; it would differentiate a different function with the same input and output shapes. The reverse rule must therefore retain the realized mask or enough state to reproduce that exact draw.

Each dropout measurement starts a fresh session and evaluates before performing the batch update. Within that session, the evaluation and the training loss therefore refer to the same initial parameter state. The training loss is produced as part of the update; it is not a second evaluation of the updated model. Evaluating after the update would mix the effect of mode with the effect of changed parameters.

Dynamic Control Flow And Typed Graphs🔗

A fixed typed model graph needs its operation structure and its shapes to be known when it is recorded. A program that reads a token value and then changes that structure cannot be represented by this recorder as one fixed graph.

This differs from a scalar conditional inside a tensor computation. The custom computation frontend supports if expressions on tensor elements on CPU and GPU. The conditional stays in the scalar program; it does not change the model graph's nodes or shapes.

Tracing and export illustrate why this restriction matters. Consider a function whose branch depends on the input values:

# A value-dependent branch must remain observable for later
# inputs.
def f(x):
    if x.sum() > 0:
        return x * 2
    else:
        return x - 100

Tracing on a positive-sum input records only the first branch. PyTorch warns that the trace may not generalize:

>>> # Trace the positive branch, then probe a negative input
>>> # with the same shape.
>>> traced = torch.jit.trace(f, torch.tensor([1.0, 1.0]))
TracerWarning: Converting a tensor to a Python boolean might cause the
trace to be incorrect. ... the trace might not generalize to other inputs!
>>> f(torch.tensor([-1.0, -1.0])).tolist()
[-101.0, -101.0]
>>> traced(torch.tensor([-1.0, -1.0])).tolist()
[-2.0, -2.0]
>>> traced.graph
graph(%x : Float(2, strides=[1], requires_grad=0, device=cpu)):
  %5 : Long(requires_grad=0, device=cpu) = prim::Constant[value={2}]()
  %6 : Float(2, strides=[1], requires_grad=0, device=cpu) = aten::mul(%x, %5)
  return (%6)

The recorded graph contains one aten::mul and no comparison. It therefore returns [-2.0, -2.0] on the negative input, where the original function takes the unrecorded else branch and returns [-101.0, -101.0]. Matching the input shape did not preserve the input-dependent choice.

The modern export path refuses instead:

>>> # Ask export to retain a branch whose decision depends
>>> # on a tensor value.
>>> class F(torch.nn.Module):
...     def forward(self, x):
...         return f(x)
>>> torch.export.export(F(), (torch.tensor([1.0, 1.0]),))
GuardOnDataDependentSymNode: Could not guard on data-dependent expression
Eq(u0, 1) (unhinted: Eq(u0, 1)).

The partial graph ends at torch.ops.aten.item.default: evaluating the branch condition requires a tensor value that is not determined by the input shape.

TorchLean's typed graph recorder supports a fixed operation structure and rejects unsupported recording operations. Its shape indices check tensor dimensions and references; they do not prove that an arbitrary frontend program preserves data-dependent branching. A fixed graph must represent the intended choice explicitly rather than specialize it to one sample's values.

Data-dependent structure can remain in an eager frontend that supports it. Other options are to express the choice as a supported tensor operation, such as a mask or select, or record separate static branches and choose between them at runtime. In each case, the implementation must preserve the choice on subsequent inputs.

Agreement on several positive-sum inputs would not expose this trace's defect: all of them follow the branch that recording retained. A useful probe crosses the decision boundary while keeping the tensor shape fixed. The negative input above does exactly that. More generally, validating a captured program requires inputs that distinguish its possible behaviors, not just repeated examples from the region in which it was captured.

A branch on a fixed model configuration can be resolved while constructing a graph. A branch on a future tensor value needs a representation of that decision in the executable program. The polymorphic TorchLean function interface uses opaque tensor references even with an eager handler; it does not expose their elements as ordinary Lean booleans. Supporting that style of dynamic control flow would require an appropriate frontend and runtime interface, beyond the current model API.

Unsupported Runtime Configurations🔗

The following requests fail at different configuration boundaries. The messages are wrapped here to fit the page; each is a single line in the terminal.

An unimplemented device:

Recorded output
$ # Request a recognized device that has no maintained
$ # runtime profile.
$ scripts/lake.sh exe torchlean quickstart_mlp --device metal --steps 1
error: quickstart_mlp: device `metal` has no
maintained runtime profile; use a programmatic backend profile

An arithmetic the lower dispatcher can execute but the supervised trainer cannot use, --arithmetic complex, is refused when the flags are parsed; Command-Line Reference shows the message.

A CUDA build with no visible device:

Recorded output
$ # Hide visible GPUs while keeping the native CUDA build to
$ # isolate runtime availability.
$ CUDA_VISIBLE_DEVICES="" scripts/lake.sh -K cuda=true exe torchlean quickstart_mlp --device cuda
error: torch eager session: CUDA was requested and
this is a CUDA build, but no usable CUDA device is visible

That last message distinguishes the two failure modes it could be reporting. A build without LibTorch says so and tells you to rebuild with -K cuda=true; a real CUDA build with no device says that instead. Both refuse, and both are raised when the session opens rather than when the flags are parsed, so session validation checks programmatic configurations as well as command-line requests.

The rejection is a value in Lean, so this page can exhibit one without a terminal. A configuration built in code passes no flag parser, so the same complex-arithmetic request fails when the session opens:

-- Ask the supervised trainer for complex arithmetic and -- catch the session-opening failure. rejected: TorchLean.Trainer: supervised training supports real arithmetic; complex arithmetic requires an explicit complex-valued training API #eval do try let trainer := Trainer.new runModel { objective := mse, arithmetic := complex, seed := 2026 } let _ ← trainer.train runData { steps := 1, logEvery := 0 } IO.println "the run was accepted" catch e => IO.println s!"rejected: {e.toString}"
rejected: TorchLean.Trainer: supervised training supports real arithmetic; complex arithmetic requires an explicit complex-valued training API

The general rule covers more than these three. CUDA in a CPU-only build, proof-only arithmetic in IO, and an operation with no admissible capsule all fail rather than quietly changing the requested configuration behind the caller's back. These failures help protect benchmark provenance. Record the accepted configuration and actual provider so that a requested device is not mistaken for evidence of where an operation ran.

Choosing A Configuration🔗

Goal

Arithmetic

Profile

inspect ordinary training

native Float32

CPU

run native GPU training

native Float32

CUDA

inspect binary32 reference behavior

FloatLib binary32

CPU

run Lean-composed attention on GPU primitives

native Float32

CUDA through LibTorch

verify or export an operation graph

semantic context

the IR semantics, no trainer profile

The final row lowers a model to NN.IR.Graph for semantic inspection and proof. See The Specification Layer for how to reason about that graph.

Reproducible Comparisons🔗

To make a run comparison reproducible, record:

model architecture and parameter count
dataset identity and preprocessing
seed and optimizer
arithmetic
device and provider capsules
train or evaluation mode
checkpoint and code revision

Without those, two loss curves may be incomparable even when both are labelled "TorchLean float32". The batching comparison has a nonzero gap that a six-decimal summary hides. Provider and reduction-policy metadata help investigate such differences; exact buffer comparisons and controlled experiments are needed to explain them.

Sources: