An Affine Model With Explicit Parameters🔗

An affine function isolates the relationship between a model's parameters and its input. The record stores a weight and bias; forward receives both the record and the input:

6.500000

The result of WfAffine.forward depends on its WfAffine and Float arguments. Fixing them determines the computation, so a proposition can refer to that same application:

The coefficients and intermediate values are exactly representable, so rfl can unfold the definitions to the same Float value on each side. The proposition names the parameter record and input directly; those arguments identify the entire computation. This proof works by reduction. It does not give Float arithmetic the algebraic laws of the reals.

Tactic `rfl` failed: The left-hand side
  0.1 + 0.2
is not definitionally equal to the right-hand side
  0.3

⊢ 0.1 + 0.2 = 0.3

55.511151

Parameter Payloads🔗

TorchLean's layer representation takes the same idea to tensor scale. A Layer σ τ records:

• stateShapes, the shapes of its parameters and persistent buffers, in the order forward expects them;

• initState, initial values for that state, as Float tensors;

• requiresGrad, one flag per state entry, so buffers can be marked as not differentiated;

• updateBuffers, an optional function for running statistics;

• forward, a program that takes a mode and then the state followed by the input.

Every one of those is a field of an ordinary structure. Training replaces the parameter values many times without touching the architecture, because the architecture is the value that holds the shapes and the forward program.

For a two-layer MLP:

[[4, 2], [4], [1, 4], [1]]

Those are the first weight matrix, the first bias, the second weight matrix, and the second bias. The order is part of the forward program's type rather than a convention to remember, and nn.initialState wfInit is the matching initial payload:

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

The state is indexed by a list of shapes, not just by the number of entries in that list. Four tensors of the wrong extents therefore cannot stand in for these four tensors. Within a layer, requiresGrad identifies which state entries are trainable, while persistent buffers can remain part of the same state interface. A forward pass may read both. This distinction is needed for layers such as batch normalization: an optimizer updates learned coefficients, whereas a buffer update records running statistics under a separate rule.

A trained runtime owns a later payload with that same ordered shape list. Models with incompatible state shapes cannot exchange payloads directly; models with identical state shapes still need a value-level check to establish which parameters they use.

Type mismatch
  nn.initialState wfWide
has type
  nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes wfWide)
but is expected to have type
  nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes wfInit)

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

17

(17, 17)

Seeded Initialization🔗

The linear layers need random initial weights. Rather than drawing them from an unnamed global generator, layer constructors return an nn.Builder, which is a deterministic computation over a seed stream. nn.build runs that computation with an initial seed:

@nn.build : {α : Type u_1} → ℕ → nn.Builder α → α

Here α is the type of value the builder produces, a sequential model in this example. nn.build takes the seed and builder and returns that value directly, without IO. Repeating the same arguments gives the same value. A custom builder could use the seed to choose an architecture too; our fixed builder uses it to draw the weights:

2026: [[2, 2]: [[-0.534110, 0.833175], [-1.213642, 0.820955]], [2]: [0.000000, 0.000000]]
   7: [[2, 2]: [[0.964935, -0.659584], [0.073548, -1.193925]], [2]: [0.000000, 0.000000]]
2026 twice equal: true

The two seeds give different weight matrices and the same zero bias. The final Boolean compares formatted state strings, so it checks only what the printer exposes. Exact reproducibility follows from applying the pure nn.build definition to the same seed and builder.

The builder's seed stream is independent of unrelated random draws. Repeating the build after drawing randomness and initializing another model gives:

unrelated build: [[8, 8], [8]]
still equal: true

# Compare initialization before, after, and with an extra
# draw following a seed reset.
import torch, torch.nn as nn

a = nn.Linear(2, 2)
b = nn.Linear(2, 2)
print(torch.equal(a.weight, b.weight))

torch.manual_seed(2026); c = nn.Linear(2, 2)
torch.manual_seed(2026); d = nn.Linear(2, 2)
print(torch.equal(c.weight, d.weight))

torch.manual_seed(2026); e = nn.Linear(2, 2)
torch.manual_seed(2026); _ = torch.randn(1); f = nn.Linear(2, 2)
print(torch.equal(e.weight, f.weight))

False
True
False

wfFromBuilder.predict : Tensor Float [2] → IO (Tensor Float [1])

wfFromValue.predict : Tensor Float [2] → IO (Tensor Float [1])

Optimizer State Transitions🔗

Let \theta_t be the parameter tensor at step t, g_t its loss gradient, and \eta the learning rate. Plain gradient descent updates the parameters by

\theta_{t+1}=\theta_t-\eta g_t,

and momentum (Polyak, 1964) adds a velocity v_t whose previous value is weighted by the coefficient \mu:

\begin{aligned} v_{t+1} &= \mu v_t+g_t,\\ \theta_{t+1} &= \theta_t-\eta v_{t+1}. \end{aligned}

The step returns both the updated parameters and the updated velocity. TorchLean packages these as parameters and optimizer state:

Optim.Step : (α : Type) → [Storage α] → Shape → Type → Type

Optim.Step α s σ pairs the next parameters, of shape s, with the next optimizer state, of type σ. The momentum optimizer is then a function from the old triple to that pair:

@Optim.MomentumSGD.update : {α : Type} →
  [inst : Storage α] →
    [inst_1 : Context α] →
      [DecidableRel fun x1 x2 => x1 > x2] →
        {s : Shape} →
          Optim.MomentumSGD.State α s → Tensor α s → Tensor α s → Optim.Step α s (Optim.MomentumSGD.State α s)

The next call needs both fields of Step. Passing new parameters with an old momentum buffer changes the recurrence. The shared shape variable aligns parameters, gradients, and buffer entries coordinate by coordinate; the storage and arithmetic instances supply their representation and operations.

Let us run two steps of it from \theta_0 = 2, with \eta = 0.1, \mu = 0.9, and a constant gradient g = 0.25. By hand: v_1 = 0.25, so \theta_1 = 2 - 0.025 = 1.975; then v_2 = 0.9\cdot 0.25 + 0.25 = 0.475, so \theta_2 = 1.975 - 0.0475 = 1.9275.

step 1: [1.975000], buffer [0.250000]
step 2: [1.927500], buffer [0.475000]

The gradient stays fixed, but the second update grows because the buffer retains the first gradient's contribution. Resuming this trajectory therefore needs the momentum buffer as well as the parameter. Saving only the parameter is enough to reuse the forward function, but loses the state needed for the next training step.

PyTorch, asked for the same two steps:

# Feed the same explicit gradient twice and inspect the
# stored momentum buffer.
p = torch.nn.Parameter(torch.tensor([2.0]))
opt = torch.optim.SGD([p], lr=0.1, momentum=0.9)
for _ in range(2):
    opt.zero_grad()
    p.grad = torch.tensor([0.25])
    opt.step()
    print(p.detach().tolist(), opt.state[p]["momentum_buffer"].tolist())

[1.975000023841858] [0.25]
[1.9275000095367432] [0.4749999940395355]

The values follow the same recurrence at the displayed precision; the PyTorch example uses binary32 and the Lean example binary64. MomentumSGD.State's docstring specifies PyTorch's SGD convention with dampening = 0 and nesterov = false. These two steps provide a numerical check of that convention.

Nothing in the Lean version mutates the old state; it returns the new one. Adam and AdamW carry more fields (Loshchilov and Hutter, 2019), and the picture is unchanged.

# Deliberately keep earlier gradients to show how
# accumulation changes the update.
p = torch.nn.Parameter(torch.tensor([2.0]))
opt = torch.optim.SGD([p], lr=0.1)
for i in range(3):
    (0.25 * p).sum().backward()   # no opt.zero_grad()
    opt.step()
    print(i, p.grad.tolist(), p.detach().tolist())

0 [0.25] [1.975000023841858]
1 [0.5]  [1.9250000715255737]
2 [0.75] [1.850000023841858]

[1.975000] [1.950000] [1.925000] 
[1.975000] [1.925000] [1.850000] 

([1.975000], [1.975000], true)

Tactic `rfl` failed: The left-hand side
  wfMomStep[0]
is not definitionally equal to the right-hand side
  wfSgdStep[0]

⊢ wfMomStep[0] = wfSgdStep[0]

Training And Evaluation Modes🔗

Some layers denote two different functions. Dropout samples a mask during training and is the identity during evaluation (Srivastava et al., 2014). Batch normalization uses batch statistics and updates running buffers during training, then reads the saved statistics during evaluation (Ioffe and Szegedy, 2015). In TorchLean, Mode is an argument of Layer.forward, ahead of the state and the input.

Runtime.Autograd.Model.Layers.Mode.train : Runtime.Autograd.Model.Layers.Mode

The running-statistics update is a function too, with the batch value and the old running value both named:

[1.200000, 2.200000]

The first vector holds the old running values and the second holds the batch statistics. The rank-zero tensor weights the new batch by one tenth. Optimizer momentum instead weights the old buffer, so the shared name hides a difference in the equations.

That is 0.9\cdot(1,2) + 0.1\cdot(3,4), the exponential moving average with momentum 0.1. The corresponding PyTorch running_mean update is:

# Use identical rows so the batch mean is exactly the vector
# being blended into the buffer.
bn = nn.BatchNorm1d(2, momentum=0.1, affine=False)
with torch.no_grad():
    bn.running_mean.copy_(torch.tensor([1.0, 2.0]))
bn.train()
bn(torch.tensor([[3.0, 4.0], [3.0, 4.0]]))
print(bn.running_mean.tolist())

[1.2000000476837158, 2.200000047683716]

For dropout, let p be the probability of dropping an entry. In evaluation the layer is the identity; in training this construction uses a seeded mask and scales the surviving entries by 1/(1-p). For ideal independent Bernoulli masks this scaling preserves each entry in expectation; that probabilistic statement is distinct from inspecting one deterministic mask:

eval  = [1.000000, 1.000000, 1.000000, 1.000000]
train = [2.000000, 0.000000, 2.000000, 0.000000]

Evaluation leaves the four ones unchanged. Training keeps two entries and doubles them. This particular mask preserves the input sum, but another mask need not: the expectation statement concerns repeated ideal draws at each coordinate.

PyTorch, same layer and same input:

# Observe evaluation identity and one training mask from
# PyTorch's generator.
torch.manual_seed(2026)
drop = nn.Dropout(0.5)
x = torch.ones(4)
drop.eval();  print(drop(x).tolist())
drop.train(); print(drop(x).tolist())

[1.0, 1.0, 1.0, 1.0]
[2.0, 0.0, 2.0, 2.0]

The convention agrees: identity in evaluation, and survivors scaled by 2 in training, which is the inverted-dropout scaling from the original paper's section 10. The masks differ because the two masks came from different generators. Ours came from the seed the builder consumed, so nn.build 2026 (nn.dropout 0.5) reproduces that seeded construction, while PyTorch's came from the global generator's state at the moment of the call.

Making the mode an argument is also what lets an exported graph or a theorem say which function it is about. A statement about a model containing dropout is ambiguous until the mode is fixed, and here it cannot be left unfixed, because forward will not run without it.

Runtime Effects And Pure Interfaces🔗

Reading a dataset and launching a kernel are effects. The call that hands a model to a runtime therefore has a different result type from the pure builder:

@nn.Module.eval : {σ τ : Shape} →
  {α : Type} → [inst : Storage α] → [inst_1 : Context α] → {model : nn.Sequential σ τ} → nn.Module α model → IO Unit

Read the final ordinary argument as the live module whose mode will change. The implicit model argument ties that module to its architecture, and σ and τ record its input and output shapes. None of those arguments is an input tensor to predict on. After the returned action runs, a subsequent forward call reads the selected mode. This explains why m.eval and #eval expression have unrelated roles despite their similar spelling: one changes runtime state, while the other asks Lean to execute an expression during the example.

The IO Unit result says that switching a live module into evaluation mode is an action with no returned data. For a Float model with output shape [1], nn.Module.forward returns IO (Tensor Float [1]): the action produces a tensor.

The same separation works for a saved certificate. IO.FS.readFile reads its bytes; a pure parser inspects the returned string; a pure checker inspects the parsed certificate. A theorem can then state what acceptance implies without modelling the file system. Training likewise performs its parameter updates, tape operations, and device activity in IO, while the specification layer reasons about pure interpretations of the model.

Storage Reuse In The Runtime🔗

Lean 4's reference counting allows storage reuse when an immutable value has a unique owner: the compiled code can update it in place (Ullrich and de Moura, 2019). Above that, TorchLean uses explicitly mutable runtime objects and foreign buffers wherever the workload calls for them; Backend Selection is where that machinery lives.

The source-level interface controls dependencies. Storage ownership and mutation belong to the execution strategy. That is also what lets several backends implement one model: a CPU evaluator, a native CUDA kernel, and an external provider all receive the same explicit inputs, so changing device does not mean rewriting the model around a different collection of hidden fields.

Function Inputs, State, And Effects🔗

When you meet a TorchLean definition, three questions place it quickly:

1. Which values determine the result?

2. Is it pure, a deterministic computation over a seed stream such as nn.Builder, or an IO action?

3. If state changes, where is the old state named and where is the new state named?

Try them on Optim.MomentumSGD.update. Its result depends on the optimizer state, the parameters, and the gradients, and on nothing else. It is pure. The old state is the first argument, and the new state is the optimizerState field of the returned Optim.Step. Now try them on autograd.model.grad, which appears in TorchLean and PyTorch: it takes a model, a loss, a state, an input, and a target, and it returns its answer in IO, because it runs a tape on a real device.

References🔗

Reference counting for a purely functional language is Ullrich and de Moura (2019). The two mode-dependent layers used as examples are Srivastava et al. (2014) and Ioffe and Szegedy (2015); momentum is Polyak (1964) and the weight-decay variant is Loshchilov and Hutter (2019). The floating-point discussion follows Goldberg (1991) and Boldo and Melquiond (2011). The framework we compare against throughout is Paszke et al. (2019).

• Lean 4 language reference: Functions and Do notation.

• PyTorch reproducibility notes: Randomness.

• PyTorch API pages: torch.optim.SGD, nn.Dropout, and nn.BatchNorm1d.

• TorchLean sources: Layer, Seq, and Optim.