1.6. TorchLean and PyTorch
TorchLean looks familiar on purpose. It has tensors, modules, parameters, autograd, optimizers, devices, and checkpoints because those ideas already work well in modern ML. But TorchLean is not PyTorch rewritten in Lean, and it is not trying to catch PyTorch by accumulating the same number of operators.
PyTorch's center of gravity is execution. A Python program can assemble a large model dynamically, train it with highly tuned kernels, distribute the work over many devices, and deploy the result. TorchLean's center of gravity is the connection between a running model and a mathematical statement. It can execute models too, but it keeps asking questions that PyTorch normally leaves to the surrounding project: What is the exact shape contract? Which parameter payload did the verifier read? What equation does this graph node denote? Which arithmetic appears in the theorem?
That difference is easier to see in code than in slogans, so we will build the same MLP in both systems and follow it through initialization, autograd, compilation, and execution.
1.6.1. The Same MLP In Both Systems
A PyTorch model commonly owns its parameters through nn.Module:
# Python / PyTorch
model = torch.nn.Sequential(
torch.nn.Linear(4, 8),
torch.nn.ReLU(),
torch.nn.Linear(8, 2),
)
logits = model(x)
The corresponding TorchLean builder is:
import NN.API
open TorchLean
def model :
nn.M (nn.Sequential (.dim 4 .scalar) (.dim 2 .scalar)) :=
nn.Sequential![
nn.linear 4 8,
nn.relu,
nn.linear 8 2
]
Both programs describe an affine map, ReLU, and another affine map. Their surrounding contracts are different.
In PyTorch, an eager tensor carries its shape as runtime metadata. Calling a layer inspects those
dimensions while the program executes. PyTorch's export and compilation systems can add symbolic
shape constraints later, but an ordinary annotation such as torch.Tensor does not distinguish a
vector of length four from a matrix with four columns.
In TorchLean, the input and output shapes index the nn.Sequential type. Layer composition is
checked while Lean elaborates the definition. nn.M also records that model is a deterministic
seed-state computation waiting to initialize its parameters.
Dynamic shapes are convenient during exploration and for data-dependent programs. Shape-indexed types require more information up front, but they make layer composition and later theorem statements much cleaner. TorchLean still accepts runtime-loaded data; it checks the dimensions once at the boundary and then works with the resulting typed tensor.
1.6.2. Parameters: Object Fields Versus Explicit Payloads
A PyTorch nn.Module registers parameter objects. Calling model(x) reads the module's current
fields. An optimizer mutates those parameters, usually through gradient fields populated by
autograd.
TorchLean's model description contains parameter shapes, initialization tensors, gradient flags, and a forward program. The forward program receives the live parameter payload explicitly. Initialization and execution are therefore related but distinguishable:
def initialized := nn.run 2026 model #check nn.paramShapes initialized #check nn.initParams initialized #check nn.forwardProgram (model := initialized)
The initial tensors are part of the initialized description. Training creates a runtime module from
them and updates the runtime's parameter state. A theorem about nn.initParams initialized is not a
theorem about the parameters after 10,000 optimizer steps.
This explicit payload pays off when the model becomes a graph. The compiler receives the graph and the exact tensors being analyzed rather than reading whatever happens to be stored in mutable fields at that moment. A PyTorch checkpoint importer therefore has a concrete job: map names, shapes, order, and layout into this payload.
1.6.3. Autograd: A Tape In Two Different Roles
PyTorch dynamically constructs an autograd graph while tensor operations run. During backward, saved tensors and derivative rules propagate vector-Jacobian products to leaves. The implementation is mature, broad, and deeply integrated with the dispatcher.
TorchLean's eager runtime also records a tape. The difference is that the codebase keeps three layers available side by side:
-
the runtime node and saved values used to compute a gradient;
-
the ideal derivative or VJP definition used in a theorem;
-
the proof that the runtime rule implements that mathematical derivative.
For a scalar loss L(\theta), reverse mode does not need to materialize the full Jacobian. Starting
with the cotangent 1, each node applies a local VJP:
\bar{x}=J_f(x)^\mathsf{T}\bar{y}.
The runtime rule is the fast implementation. The ideal VJP is the equation we want it to implement. The theorem, when available for that operation, connects the two.
This distinction becomes especially important for external kernels. The maintained LibTorch-forward attention path asks LibTorch to compute the forward value, records the ordinary TorchLean tape node, and uses TorchLean's local backward rule. Handing both forward and backward to LibTorch would instead trust LibTorch autograd, saved-tensor conventions, gradient extraction, and parameter ownership. Kernel capsules record which choice was made.
1.6.4. Training Loops
A conventional PyTorch loop is explicit Python mutation:
for x, y in loader:
optimizer.zero_grad()
prediction = model(x)
loss = loss_fn(prediction, y)
loss.backward()
optimizer.step()
TorchLean's trainer packages the same lifecycle:
def trainer :=
Trainer.new model
{ task := .classification
optimizer := optim.adam { lr := 0.001 }
seed := 2026 }
-- let result ← trainer.train dataset
-- { steps := 200, batchSize := 16, logEvery := 25 }
Internally, the runtime owns mutable parameters, gradients, optimizer moments, tape state, and possibly device buffers. TorchLean does not force a large GPU training loop to allocate a new pure tensor tree at every step. The semantic interfaces remain explicit while the execution engine uses mutation and ownership where performance requires it.
The train call returns a trained result whose prediction closures refer to the trained runner.
Lower-level manual APIs expose parameter tensors and individual forward, backward, and optimizer
steps when verification or research code needs them.
This is still an ordinary training loop: it produces updated parameters, logs, and predictions. The later proof chapters start from those concrete objects rather than trying to read a theorem out of the loss curve.
1.6.5. Dispatch And Backend Selection
PyTorch routes an operation through its dispatcher. Device, dtype, layout, compilation state, and available libraries determine the eventual implementation. A CUDA matrix multiplication may use cuBLAS, convolution may use cuDNN, and attention may select one of several fused kernels.
TorchLean's backend framework expresses a smaller but more explicit decision:
-
Devicesays where execution should occur; -
Providersays which implementation family supplies the operation; -
BackendOpidentifies the requested operation; -
a backend profile supplies provider preference, assurance policy, and available capsule modules;
-
a kernel capsule records shape/layout requirements, forward and VJP ownership, and numerical policy.
This planner answers a question that is often surprisingly hard to answer after a large run: which implementation actually handled each expensive operation? If LibTorch supplies attention while native CUDA supplies matrix multiplication, the audit report names both. If the requested provider is unavailable, planning fails instead of quietly choosing a different story.
This is also how TorchLean scales. It can keep ownership of the model, parameter layout, graph, and tape while calling industrial kernels for the expensive arithmetic.
1.6.6. LibTorch, ATen, And The Actual Kernel
These names are easy to mix up. LibTorch is PyTorch's C++ distribution. ATen is the tensor and operator layer used inside PyTorch. Beneath an ATen operation there may still be another library: cuBLAS for matrix multiplication, cuDNN for convolution, or a fused attention implementation.
Calling LibTorch sends a particular operation across an FFI boundary while TorchLean retains the model. The maintained scaled-dot-product-attention path works like this:
-
TorchLean owns the model and current parameter tensors.
-
TorchLean asks LibTorch for the attention forward value.
-
TorchLean stores that value and records its ordinary attention tape node.
-
During backward, TorchLean applies its own attention VJP.
Another profile can use native TorchLean CUDA for both forward and backward. A future profile could delegate both directions, but that would be a different capsule because it would trust LibTorch's autograd state as well as its forward kernel.
1.6.7. Graphs And Compilation
PyTorch offers FX, torch.export, AOTAutograd, and compiler stacks such as torch.compile.
Their graphs support transformation and deployment inside the PyTorch ecosystem.
TorchLean has two graph-facing layers with different purposes:
-
GraphSpecdescribes structured architectures and compiles them to TorchLean programs; -
NN.IR.Graphis the canonical operation DAG used by lower-level evaluation and verification.
An IR node contains an operation tag, parent ids, and an output shape. Parameter and constant values
live in payload stores. NN.IR.Semantics defines how supported nodes are interpreted over a scalar
domain.
The verification compiler can lower supported initialized models and parameter payloads to
this IR. A separate first-order source language under NN.Verification.TorchLean.Proved has an
end-to-end compiler-correctness theorem. They share an IR target, but the theorem applies to the
proved source fragment, not automatically to every model accepted by the broader executable
compiler.
This is a deliberate difference from a marketing-style "compiled" flag. TorchLean keeps the executable compiler and the proved compiler fragment separately named so users can see which guarantee they actually have.
1.6.8. Checkpoints And Graph Import
PyTorch checkpoints are Python-oriented zip/pickle artifacts. TorchLean does not duplicate that
loader inside Lean. Its adapter generates Python that asks PyTorch to load a state_dict and emit
named tensor data as plain JSON. Lean then parses each requested tensor into a statically known
shape.
Graph import is separate. A generated adapter uses torch.export or FX to emit the
torchlean.ir.v1 format, which Lean parses into NN.IR.Graph. The ONNX adapter lowers supported
static nodes to the same format.
These steps establish progressively stronger but still limited facts:
-
JSON parsing establishes that the bytes match the expected schema.
-
Tensor parsing establishes that a requested payload has the expected shape.
-
Graph checks establish node references and shape contracts.
-
Operator-by-operator semantic refinement would establish equality with the source graph.
-
A runtime bridge would connect that graph semantics to the deployed kernels.
The first three steps are already useful: they catch malformed exports and shape/layout mistakes. The remaining two are where a claim about the imported program becomes a claim about its meaning.
1.6.9. Floating Point
PyTorch's numerical behavior depends on dtype, device, library versions, compiler transformations, and reduction algorithms. Its numerical-accuracy documentation explicitly warns that mathematically identical computations are not guaranteed to be bitwise identical across batched, sliced, device, or backend paths.
TorchLean names several numerical meanings:
-
real-valued specifications for ideal mathematics;
-
NF, a configurable rounded-real arithmetic; -
FP32, a rounded-real specialization with binary32 precision and gradual underflow, but without an upper exponent bound or IEEE special values; -
IEEE32Exec, an executable bit-level binary32 model; -
runtime CPU, CUDA, and LibTorch representations.
The generic layer was influenced by Flocq's separation of formats from rounding operators.
IEEE32Exec is TorchLean's executable Lean reference for binary32. The floating-point chapters
derive these layers from examples and show how they reconnect to a runtime.
Training can proceed without this extra structure. Claims that compare an exact equation with an executable result need it.
1.6.10. Mixed Precision Is A Current Difference
PyTorch tensors carry dtype at runtime, and operators define promotion, casting, and autocast behavior. It is therefore natural for one PyTorch graph to combine integer tokens, FP8 or FP16 activations, BF16 weights, and FP32 accumulation.
TorchLean's numeric tensor type is Tensor.T α s. A current model execution and an
NN.IR.Semantics evaluation select one numeric α for parameters, activations, intermediates, and
outputs. The same model can be instantiated under another α, but two numeric nodes do not carry
independent dtype indices. Boolean masks and integer indexing data use dedicated interfaces; they
do not amount to general numeric dtype promotion.
Consequently, choosing executable binary32 or complex binary32 changes the scalar interpretation of a whole supported run. It does not reproduce PyTorch autocast or mixed-precision master weights. A faithful mixed FP8/BF16/FP16/FP32 graph will require explicit dtype-indexed values, cast and quantize nodes, promotion rules, accumulation dtypes, and backend contracts for each conversion. Until then, numerical claims should describe the homogeneous scalar semantics that actually ran.
1.6.11. What TorchLean Adds
For the small MLP, PyTorch gives an excellent route from Python source to fast training. TorchLean adds a route from a typed model to several inspectable semantic objects:
shape-typed model -> initialized parameter layout -> runtime tape and trained payload -> canonical operation graph -> exact or rounded scalar semantics -> bound or certificate -> theorem with named assumptions
Each arrow is an interface that can be tested, checked, or proved independently as coverage grows.
1.6.12. When To Use Which
PyTorch is the natural choice when the main requirement is broad model coverage, pretrained ecosystems, distributed training, compilation, and production deployment. TorchLean becomes useful when the project also needs:
-
dimensions checked in model and tensor types;
-
mathematical operations available as Lean definitions;
-
an inspectable graph and parameter payload;
-
executable verification and certificate checking;
-
formal statements about graph, autograd, optimizer, or numerical behavior;
-
an explicit ledger of native and external assumptions.
Many projects should use both. Training can remain in PyTorch while TorchLean checks an exported artifact, or training can run in TorchLean while selected bottlenecks use PyTorch kernels. They solve different parts of the same problem.
1.6.13. References
-
Adam Paszke et al., “PyTorch: An Imperative Style, High-Performance Deep Learning Library”, NeurIPS 2019.
-
Sylvie Boldo and Guillaume Melquiond, “Flocq: A Unified Library for Proving Floating-Point Algorithms in Coq”, IEEE ARITH 2011.