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 |
| which binary32 semantics every scalar operation obeys |
device and providers |
| which implementation of each operation is selected, and what evidence backs it |
train or evaluation mode |
| 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.
#print 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.
#print Runtime.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
#check @Config.withDevice
open Runtime.Autograd.Torch in
#check @Config.validateForExecution
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
#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}"
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:
# 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:
# 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:
# 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:
== 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 |
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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:
# 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 |
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
# 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.
#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}"
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 }
#eval do
let runs := [runCpu, runIeeeCpu, runCuda]
runs.forM fun (run : Trainer.RunConfig) => do
IO.println
s!"{run.toConfig.deviceName} {run.arithmetic.cliName}"
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.
#check @nn.lowerToTypedGraph
#check @nn.TypedGraphModel.forward
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.
#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] 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 |
|---|---|
| training |
| training |
| evaluation |
| evaluation |
| 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 }
#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}"
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:
$ # 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:
$ # 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.
#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}"
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 | CPU |
run native GPU training |
native | CUDA |
inspect binary32 reference behavior | FloatLib binary32 | CPU |
run Lean-composed attention on GPU primitives |
native | 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:
-
Trainer/Core.lean, the
RunConfigrecord; -
Session.lean, the single session choke point;
-
Types.lean,
validateForExecution; -
Session.lean, capsule reporting;
-
Report.lean, the evidence lines;
-
SimpleMlpTrain.lean, the runnable example.