MLP Definitions🔗

A PyTorch model commonly owns its parameters through nn.Module:

# Build a four-feature, two-output MLP and call it on one
# input vector.
# Python / PyTorch
import torch
net = torch.nn.Sequential(
    torch.nn.Linear(4, 8),
    torch.nn.ReLU(),
    torch.nn.Linear(8, 2),
)

x = torch.zeros(4)
logits = net(x)

The corresponding TorchLean builder is:

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.Builder also records that tpModel is a deterministic seed-state computation waiting to initialize its parameters.

Application type mismatch: The argument
  bc✝
has type
  nn.Sequential (Shape.appendDim [] 7) (Shape.appendDim [] 2)
but is expected to have type
  ?m.67 a✝ __r✝¹ bc✝ __r✝ (Shape.appendDim [] 8) [2]
in the application
  nn.compose a✝ bc✝

# Python / PyTorch
>>> net(torch.zeros(3))
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x3 and 4x8)

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:

nn.initialState tpInit : nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes tpInit)

The initialized description exposes its state shape list:

[[8, 4], [8], [2, 8], [2]]

The four groups contain 32 weights, eight hidden biases, sixteen readout weights, and two readout biases. The input and output interface [4] -> [2] alone does not reveal those groups: another model could use the same interface with a different hidden width. The shape list exposes that internal layout. The PyTorch names additionally identify where each tensor came from in the module tree. An adapter must connect names to positions and validate their shapes before it can compare or execute the resulting payload.

PyTorch answers the same question through the module's state_dict:

# Inspect the named state tensors that must be matched by
# any checkpoint adapter.
# Python / PyTorch
for name, t in net.state_dict().items():
    print(f"{name:10s} {tuple(t.shape)}")

0.weight   (8, 4)
0.bias     (8,)
2.weight   (2, 8)
2.bias     (2,)

Both lists contain the same four tensors in the same order. PyTorch attaches a string key to each tensor; loading must account for missing, unexpected, or renamed keys. TorchLean attaches a position and shape computed from the model. This rejects incompatible state shape lists, but another architecture with the same state shapes can supply a payload of the same type. Neither convention alone proves that the intended weights were loaded.

The initial tensors are part of the initialized description. Execution instantiates a mutable nn.Module from them, and the whole lifecycle is ordinary IO:

[-0.284787, -0.230906]

The module owns the live parameters, persistent buffers, and train/eval mode, and module.state reads the current payload back out. The immutable nn.Sequential remains the value used by graph lowering and proofs. A theorem about nn.initialState tpInit is therefore not a theorem about the parameters after 10,000 optimizer steps; a claim about trained weights must name the saved payload it analyzes.

Lowering receives the initialized model and the tensors to analyze. This fixes the values in the resulting graph artifact independently of later module updates. A PyTorch checkpoint importer must map the external names, shapes, order, and layout into this payload.

The payload distinction also applies to layers beyond this MLP. For integer-indexed embeddings, nn.embedding vocab embedDim constructs a trainable table, nn.Embedding.fromWeight weight preserves an existing table exactly, and freeze := true keeps that table in module state without requesting a parameter gradient. If the index tensor has dimensions dims, its output has dimensions dims ++ [embedDim]; repeated ids accumulate into the same gradient row.

Autograd Implementations And Specifications🔗

PyTorch dynamically constructs an autograd graph while tensor operations run. During backward, saved tensors and derivative rules propagate vector-Jacobian products to leaves through its dispatcher and autograd machinery ((Paszke et al., 2019; Baydin et al., 2018)).

TorchLean's eager runtime also records a tape. To reason about a backward step, its proofs refer to an ideal derivative or VJP definition. For covered operations, correctness theorems show that this rule is the mathematical derivative; numerical results relate it to the rule evaluated with the runtime node's saved values.

For a scalar loss L(\theta), reverse mode starts with derivative 1 at the loss and propagates derivatives backward. At a node y=f(x), let J_f(x) be the Jacobian and \bar{y} the derivative of the loss with respect to the node's output, represented as a column vector. The corresponding input derivative \bar{x} is obtained by the local VJP:

\bar{x}=J_f(x)^\mathsf{T}\bar{y}.

Each node computes this product without materializing the full network Jacobian. The ideal VJP states the intended rule; correctness and numerical error results describe its relationship to the runtime implementation under their hypotheses.

External kernels make the three layers particularly important because forward and backward may cross different boundaries. The backend chapter records that ownership per operation instead of letting it hide inside the training loop.

[-0.925000]

loss  = 2.030625
grads = [[2, 2]:
 [[0.000000, 0.000000], [2.137500, 4.275000]],
 [2]:
 [0.000000, 2.137500],
 [1, 2]:
 [[0.000000, -3.705000]],
 [1]:
 [-2.850000]]

# Copy all four parameter tensors before comparing the loss
# and parameter gradients.
# Python / PyTorch
grad_net = torch.nn.Sequential(
    torch.nn.Linear(2, 2), torch.nn.ReLU(), torch.nn.Linear(2, 1))
with torch.no_grad():
    grad_net[0].weight.copy_(torch.tensor([[0.5, -0.5], [1.0, 0.25]]))
    grad_net[0].bias.copy_(torch.tensor([0.1, -0.2]))
    grad_net[2].weight.copy_(torch.tensor([[1.5, -0.75]]))
    grad_net[2].bias.copy_(torch.tensor([0.05]))

loss = torch.nn.functional.mse_loss(grad_net(torch.tensor([1.0, 2.0])),
                                    torch.tensor([0.5]))
loss.backward()

loss = 2.0306248664855957
0.weight grad = [0.0, 0.0, 2.1374998092651367, 4.274999618530273]
0.bias grad = [0.0, 2.1374998092651367]
2.weight grad = [0.0, -3.7049996852874756]
2.bias grad = [-2.8499999046325684]

#[0, 0, 1074318540, 1082707148, 0, 1074318540, 0, 3228376759, 3224790630]

# Reinterpret gradient storage as 32-bit words; do not
# numerically convert the values to ints.
# Python / PyTorch
bits = []
for _, p in grad_net.named_parameters():
    bits += p.grad.view(torch.int32).flatten().tolist()
print([b & 0xffffffff for b in bits])

[0, 0, 1074318540, 1082707148, 0, 1074318540, 0, 3228376759, 3224790630]

1074318541

-190.734863

Training Loops🔗

A conventional PyTorch loop is explicit Python mutation:

# Clear old gradients before differentiating this batch and
# updating the parameters.
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:

The output is a vector of two class logits, so its class axis is zero. In oneHotCrossEntropy 0, that argument selects the axis; the one-hot target of shape [2] selects the desired class. The objective reduces the logits to a scalar loss. tpTrainer configures this objective and Adam; a subsequent training call supplies the dataset and performs the updates.

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. The running example runs 200 steps and compares the results with a PyTorch loop, explaining the different batch sizes and initializations.

[0.990000, -2.030000, 0.520000]

[0.943900, -2.168300, 0.612200]

# Reuse one fixed gradient for three momentum steps in
# binary64 arithmetic.
# Python / PyTorch
import torch
p = torch.tensor([1.0, -2.0, 0.5], dtype=torch.float64,
                 requires_grad=True)
g = torch.tensor([0.1, 0.3, -0.2], dtype=torch.float64)
opt = torch.optim.SGD([p], lr=0.1, momentum=0.9)
for _ in range(3):
    p.grad = g.clone()
    opt.step()
print(p.detach().tolist())

[0.9439, -2.1683, 0.6122000000000001]

true

[0.900000, -2.100000, 0.600000]

[0.700000, -2.300000, 0.800000]

# Hold the gradient fixed to isolate Adam's bias correction
# and epsilon placement.
# Python / PyTorch
def adam_run(steps):
    p = torch.tensor([1.0, -2.0, 0.5],
                     dtype=torch.float64,
                     requires_grad=True)
    g = torch.tensor([0.1, 0.3, -0.2],
                     dtype=torch.float64)
    opt = torch.optim.Adam([p], lr=0.1,
                           betas=(0.9, 0.999), eps=1e-8)
    for _ in range(steps):
        p.grad = g.clone()
        opt.step()
    return p.detach().tolist()

print(adam_run(1))
print(adam_run(3))

[0.900000009999999, -2.099999996666667, 0.5999999950000002]
[0.7000000299999976, -2.29999999, 0.799999985]

9.999999

29.999998

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 records the corresponding choice as a device, operation, provider, and accepted kernel capsule. Reports name the implementation selected for each operation, and unavailable requests fail instead of becoming an unreported CPU run. This keeps the model, parameter layout, graph, and tape stable while selected operations cross a native boundary.

Backend Selection explains capsules, provider preference, LibTorch forward ownership, and assurance policies in full. GPU and CUDA then follows the native CUDA boundary.

>>> # Inspect CPU and CUDA dispatch keys, then count the
>>> # recorded linear registrations.
>>> torch._C._dispatch_key_set(torch.zeros(4))
DispatchKeySet(CPU, ADInplaceOrView, AutogradCPU, AutocastCPU)
>>> torch._C._dispatch_key_set(torch.zeros(4, device='cuda'))
DispatchKeySet(CUDA, ADInplaceOrView, AutogradCUDA, AutocastCUDA)
>>> d = torch._C._dispatch_dump("aten::linear")
>>> len([l for l in d.splitlines() if ': registered at' in l])
15

checked_cpu: 6 capsules
  reference.reshape
  reference.permute
  reference.matmul
  reference.broadcast
  reference.add
  reference.relu
checked_cuda: 6 capsules
  native_cuda.reshape
  native_cuda.permute
  native_cuda.matmul
  native_cuda.broadcast
  native_cuda.add
  native_cuda.relu
cuda_request_cpu_build: rejected
  no admissible kernel capsule for op reshape on device cuda

>>> # The layer retains CPU parameters while its input is
>>> # allocated on CUDA.
>>> torch.nn.Linear(4, 2)(torch.zeros(4, device='cuda'))
RuntimeError: Expected all tensors to be on the same device, but got mat1 is
on cuda:0, different from other tensors on cpu (when checking argument in
method wrapper_CUDA_addmm)

Graphs, Lowering, 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:

• GraphSpec describes structured architectures and lowers them to TorchLean programs;

• NN.IR.Graph is 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 two graphs for our four-input MLP are a good way to see the difference in granularity. torch.export produces nine nodes:

# Export the existing model and inspect how this graph
# represents its two affine layers.
# Python / PyTorch
exported = torch.export.export(net, (torch.zeros(4),))
for node in exported.graph_module.graph.nodes:
    print(f"  {node.op:15s} {node.target}")

placeholder     p_0_weight
placeholder     p_0_bias
placeholder     p_2_weight
placeholder     p_2_bias
placeholder     input
call_function   aten.linear.default
call_function   aten.relu.default
call_function   aten.linear.default
output          output

TorchLean's verification lowering produces eighteen:

Except.ok (18, 0, 17)

aten.linear.default retains the affine layer as one graph operation; its execution can depend on later dispatch and compilation choices. TorchLean's verification graph expands that layer into reshapes, a matrix product, a broadcast, and an addition. Each primitive has a semantic equation used when reasoning through the graph one node at a time. The graph chapter lists the node kinds.

The counts describe this difference in granularity. Eighteen nodes versus nine gives no direct comparison of memory or speed; those depend on how the operations execute.

The verification lowering pass can lower supported initialized models and parameter payloads to this IR. A separate first-order source language under NN.Verification.Builtin.Proved has an end-to-end lowering-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 lowering pass.

TorchLean calls the conversion to NN.IR.Graph verification lowering. Typed graph lowering is a separate path: it records a shape-indexed SSA graph whose nodes store forward and differentiation functions. Derivative correctness is established by separate theorems over proof-carrying nodes and graphs. Neither lowering path performs optimization, fusion, scheduling, or native code generation, so selecting typed graph execution does not imply a torch.compile-style compiler.

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.

Parsing checks the JSON schema and requested tensor shapes; graph checks validate node references and shape contracts. These catch a malformed export, but a transposed square weight matrix can pass them. A value comparison can expose that error. To prove equivalence with the source graph, we would need refinement for each operator and a bridge from graph semantics to the deployed kernels. The round trip chapter exercises the import and comparison on a real checkpoint.

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.

Comparing a batched product with the product of one selected slice exposes this issue:

# Compare a row computed as part of a batch with the same
# row computed on its own.
# Python / PyTorch
a, b = torch.randn(64, 128), torch.randn(128, 128)
full, sliced = a @ b, a[:1] @ b
print(torch.equal(full[:1], sliced))
print((full[:1] - sliced).abs().max().item())
print(full[0, 0].view(torch.int32).item() & 0xffffffff,
      sliced[0, 0].view(torch.int32).item() & 0xffffffff)

False
8.58306884765625e-06
3206383400 3206383430

False says at least one entry differs; the next line gives the maximum absolute difference across the row. The integer words describe only the first entry, which need not attain that maximum. The snippet does not fix a random seed, so these are results from the recorded run.

On the recorded platform, these two mathematically equal products differed in their computed binary32 values. Blocking and reduction choices can change rounding; the displayed first entry differs by thirty ULPs. Other platforms or inputs may give equal results. A claim about the executed output may therefore need to include the batching policy: grouping requests can change the numerical computation even when the ideal function acts independently on each row.

TorchLean names several numerical meanings so that a statement can pick one:

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

• FloatLib binary32, an executable bit-level binary32 model;

• runtime CPU, CUDA, and LibTorch representations.

For batching, the library states a property of the reference semantics. If a batched run is defined by applying the same function independently to each row, then reading one row back out is that row's own result:

@unstack_mapLeading : ∀ {α : Type} [inst : Storage α] {batch : ℕ} {inShape outShape : Shape}
  (f : Tensor α inShape → Tensor α outShape) (xs : Tensor α (inShape.prependDim batch)) (i : Fin batch),
  (mapLeading [batch] f xs).unstack i = f (xs.unstack i)

The quantifiers range over an element type, a batch length, input and output shapes, a function f, the batch tensor xs, and a valid row index i. The left side applies f independently through mapLeading and then selects row i; the right side first selects that row and applies f once. Equality says those two constructions agree for every such argument. The bounded index Fin batch ensures the selected row exists. The hypothesis is built into the construction: each row is processed by the same pure function, without consulting the other rows.

This theorem establishes batch invariance for the reference semantics in BatchInvariance.lean. Applying it to a cuBLAS or other native kernel requires a contract relating that kernel to the specification. The named property states exactly what such a contract would need to preserve.

The generic layer was influenced by Flocq's separation of formats from rounding operators ((Boldo and Melquiond, 2011)), and the CompCert float work ((Boldo et al., 2015)) addresses connections between floating-point semantics and compiled code. FloatLib supplies executable binary formats, with binary32 as the familiar example here. The floating-point chapters derive these layers from examples and show how they reconnect to a runtime.

Verification Artifacts🔗

For this MLP, TorchLean exposes a sequence of objects connecting the model to a verification claim:

shape-typed model
  -> initialized parameter layout
  -> runtime tape and trained payload
  -> canonical operation graph
  -> exact or rounded arithmetic
  -> bound or certificate
  -> theorem with named assumptions

Each arrow is an interface that can be tested, checked, or proved independently as coverage grows.

PyTorch Integration🔗

Training can remain in PyTorch while TorchLean checks a supported exported artifact, or run in TorchLean while selected operations use PyTorch kernels. In either arrangement, a claim about the executed model must identify the imported weights, graph semantics, and kernel contracts.

References🔗

The system compared against throughout this chapter is described in Paszke et al. (2019), and the reverse-mode machinery both libraries implement is surveyed in Baydin et al. (2018). The floating-point design references are Boldo and Melquiond (2011) and Boldo et al. (2015). The PyTorch documentation pages used for the transcripts above are:

• nn.Module;

• nn.Embedding;

• torch.compile;

• autograd mechanics;

• torch.export;

• numerical accuracy notes.