The Source Program🔗

The model is a two-layer MLP:

Read the type before the body. nn.Builder says that model construction consumes a deterministic seed stream. nn.Sequential [2] [1] says that the initialized model will map a length-two tensor to a length-one tensor. The hidden width 8 is checked through composition: the first linear layer produces length eight, ReLU preserves the shape, and the second linear layer expects length eight.

Changing the second layer's input width to 7 breaks that composition:

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

The generated names come from the nn.Sequential! macro's binders. The dimensions locate the error: Shape.appendDim [] 7 appears where Shape.appendDim [] 8 was required. Lean rejects the composition before a forward pass runs.

The model returns a one-element tensor rather than a scalar because the shape-indexed layer API treats the output feature dimension uniformly. A dataset target must therefore have shape [1], not shape [] and not shape [batch].

Seeded Initialization🔗

Linear layers need initial weights and biases. Hiding randomness in a global generator would make a model definition depend on ambient state. TorchLean instead represents initialization as a pure state computation:

Running the same builder with the same seed produces the same initialization stream. The resulting reInit records both the architecture and its initial parameter values. Its parameter shapes are:

[[8, 2], [8], [1, 8], [1]]

Four tensors, in layer order: first-layer weights [8, 2], first-layer bias [8], second-layer weights [1, 8], second-layer bias [1]. The ReLU contributes nothing because it owns no parameters. The initial values themselves live in a payload of exactly that shape list:

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

The shape expression printed after nn.State Float depends on reInit, so it includes hidden state layout as well as the public input and output shapes. The first weight tensor contributes sixteen entries, its bias eight, the second weight eight, and its bias one. That is thirty-three numbers grouped into four tensors. A function consuming this state can pattern-match those groups without looking up string keys. It still needs the correct numerical payload: the type cannot identify which training run produced two states of the same layout.

The order is part of the forward-program interface. It is not a PyTorch-style dictionary of names. Checkpoint adapters may use names at an external boundary, but they must eventually construct this typed ordered payload.

>>> # Compare state names and shapes; the parameterless ReLU
>>> # contributes no dictionary entries.
>>> model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 1))
>>> list(model.state_dict().keys())
['0.weight', '0.bias', '2.weight', '2.bias']
>>> [tuple(v.shape) for v in model.state_dict().values()]
[(8, 2), (8,), (1, 8), (1,)]

The Dataset🔗

The target in the source example is:

Evaluating three corners and a held-out point shows both sloped and flat parts of the target:

[1.800000]

[-0.600000]

[0.200000]

[0.200000]

At [1, 1], the first ReLU receives two and the second receives zero, giving 0.8 * 2 + 0.2 = 1.8. At [-1, 1], those roles switch, giving -0.4 * 2 + 0.2 = -0.6. At the remaining two points both terms vanish. The target therefore contains more structure than a single affine map: a learner must discover where each contribution turns on. The held-out input is useful because a small training loss on the grid does not determine what the fitted function does between its rows.

Both ReLUs are off when x_1+x_2\le0 and x_2\le x_1, so that region is flat at 0.2. The held-out point [0.25, -0.75] lies inside this region and between the grid coordinates used for training. Its prediction tests the learned function away from the sampled inputs.

The example samples five positions on each input axis, giving 25 examples:

Data.Synthetic.squareGrid (-1.0) 1.0 5 : Tensor Float [5 * 5, 2]

The first dimension retains the expression 5 * 5 from the generator's type. Lean can reduce that expression to 25, so [5 * 5, 2] and [25, 2] are definitionally equal shapes.

Tensor.mapLeading applies the target to each row without fixing the operation to vectors or regression:

mapLeading [5 * 5] treats the first axis as a collection of samples and passes each remaining length-two tensor to reTarget. Since that function returns shape [1], the resulting target tensor has shape [25, 1]. Data.fromTensors then pairs rows with the same sample index. No reshuffling or independent target generation occurs in this construction, so the pairing is visible directly in the expression. This matters when debugging training: correct shapes alone would still permit labels to be attached to the wrong input rows.

Training🔗

The trainer combines the model, objective, optimizer, seed, and runtime choices:

-- Preserve the command-line runtime choices while setting
-- the training objective and optimizer.
def trainer := Trainer.new model
  { flags.runtime with
      objective := .meanSquaredError
      optimizer := optim.adam { learningRate := 0.03 }
      seed := flags.seed }

Run the complete checked-in command from the repository root:

# Train the checked-in MLP example for 200 updates from
# initialization seed 2026.
lake exe torchlean quickstart_mlp \
  --device cpu \
  --steps 200 \
  --seed 2026

On the current implementation, this deterministic run reports:

== Quickstart: simple MLP training (seed=2026, steps=200) ==
target(heldout)    = [0.200000]
untrained(heldout) = [-0.088261]
dataset size = 25
mean_loss(before training) = 0.495227
step 0: loss=0.250488
step 25: loss=0.586318
step 50: loss=0.847444
step 75: loss=0.003933
step 100: loss=0.023397
step 125: loss=0.069799
step 150: loss=0.000061
step 175: loss=0.029500
mean_loss(after training) = 0.002402
steps=200 arithmetic=native scalar=Float32 loss=0.495227 -> 0.002402
trained(heldout) = [0.228325]

The mean loss over the dataset falls from 0.495227 to 0.002402, a factor of roughly two hundred. The per-step losses fluctuate much more: 0.847444 at step 50 and 0.000061 at step 150. Each of those values measures the single example drawn at that step, so they do not trace the dataset's mean loss. After 200 such updates, the held-out prediction is 0.228325 against a target of 0.2. The run has reduced the error substantially while leaving a measurable gap.

Whenever we return to quickstart_mlp, we keep this seed and configuration fixed; later chapters also introduce smaller purpose-built variants.

$ lake exe torchlean quickstart_mlp --device cpu --steps 200 \
    --seed 2026 --arithmetic ieee
...
mean_loss(after training) = 0.002402
steps=200 arithmetic=ieee scalar=IEEE32Exec loss=0.495227 -> 0.002402
trained(heldout) = [0.228325]

# Compare a full-batch PyTorch run; its initial weights and
# update schedule differ.
torch.manual_seed(2026)
model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 1))
opt = torch.optim.Adam(model.parameters(), lr=0.03)
for step in range(200):
    opt.zero_grad()
    loss = nn.MSELoss()(model(grid), ys)
    loss.backward()
    opt.step()

target(heldout)    = 0.20000000298023224
untrained(heldout) = -0.185332
mean_loss(before)  = 0.510072
mean_loss(after)   = 1.7e-05
trained(heldout)   = 0.199078

Training Step🔗

For one sampled pair (x,y), the runtime performs:

1. read the current four parameter tensors;

2. compute Linear -> ReLU -> Linear;

3. compute the regression loss;

4. seed the loss cotangent with one;

5. traverse the autograd tape backward to obtain parameter gradients;

6. update Adam's first and second moments;

7. replace each parameter with its updated value;

8. release or retain runtime buffers according to ownership.

Write the four parameter tensors as \theta=(W_1,b_1,W_2,b_2). The weight matrices map two input coordinates to eight hidden coordinates, then eight hidden coordinates to one output; the biases have lengths eight and one. The forward map is

f_\theta(x) =W_2\,\operatorname{ReLU}(W_1x+b_1)+b_2.

Step two has a specification-level twin we can run directly. Examples.mlpForward is that equation as a pure function of two weight-and-bias records, without an autograd tape. Its tensor operations can still allocate intermediate buffers:

[-0.088261]

The pattern match reads the four tensors from the initial state. Its prediction agrees with untrained(heldout) = [-0.088261] in the training log at every printed digit. Both paths receive the same parameters; only the eager runtime records a tape. If they disagreed, this small forward calculation would help separate an error in the MLP equations from an error in loading or executing the state. It also gives proofs a computation to refer to without including tape management. The comparison checks one input; equality for all inputs requires a theorem. The derivative and interval examples below continue to use this initial payload.

The loss for a single example is a scalar function of \theta, even though the prediction has shape [1]. Reverse mode computes vector-Jacobian products from that scalar back to all four tensors (Baydin et al., 2018). Adam then uses those gradients and its optimizer state to construct the next payload.

To reason about that update, we need to connect the eager runtime's saved tensors and tape nodes to an ideal VJP. The autograd proofs state the derivative rule; rounded VJP error transformers bound the numerical discrepancy. Optimizer contracts then describe how those gradients enter SGD, momentum, or AdamW updates (Loshchilov and Hutter, 2019).

Exact Representation Of The Target🔗

The remaining training error does not come from an inability to represent the target. Its definition supplies parameters for a two-unit network directly. The term 0.8\operatorname{ReLU}(x_1+x_2) is one hidden unit with incoming row (1,1) and outgoing weight 0.8; -0.4\operatorname{ReLU}(x_2-x_1) is a second with row (-1,1) and outgoing weight -0.4; the +0.2 is the output bias. Setting those values explicitly gives:

([0.200000], [1.800000])

0.000000

true

The hidden rows compute the two linear forms in reTarget, and the readout applies their coefficients. We obtained these weights directly from the target equation. The final Boolean confirms that their computed grid loss is zero, even beyond the six printed decimals, and their held-out prediction is 0.2.

Over the reals, the construction represents the target by definition. The floating-point test covers this finite grid; it does not establish equality for every input, including overflow and non-finite cases. Still, insufficient width cannot explain the training log's remaining error: two hundred single-example Adam steps did not find a function that the architecture can represent. That distinction separates approximation theory from optimization theory. Even a zero-loss run need not recover these particular weights, as the following symmetries show.

(0.000000, 0.000000)

(true, true)

Checking The Gradient Against A Difference Quotient🔗

Decreasing loss is not enough to validate a backward pass: some incorrect updates can still reduce the objective. A separate derivative check compares the backward implementation with changes in the forward computation (Griewank and Walther, 2000).

We can estimate a derivative by perturbing a parameter and taking a difference quotient, as torch.autograd.gradcheck does. Here Examples.mlpBackward returns five tensors in order: gradients for the first weights, first bias, second weights, second bias, and finally the input. To compare them with the forward computation, return to the seeded eight-unit model and repackage its payload as two weight-and-bias records. Its first-layer pre-activations at the held-out point are:

[false, false, false, false, false, true, false, false]

(-0.479660, 0.154824)

Only hidden unit five has a positive pre-activation at [0.25, -0.75]; the other seven output zero after ReLU. This is the activation pattern of the randomly initialized model, separate from the two-unit construction of the target. It lets us inspect the gradient through one active path.

Hold the other parameters fixed and view the loss as a function of the first-layer weights. At the held-out point, let y be the prediction and t the target. The loss is their squared difference, so the cotangent supplied to the backward pass is \partial L/\partial y = 2(y-t):

0.083094

[[0.000000, -0.000000],
 [-0.000000, 0.000000],
 [-0.000000, 0.000000],
 [0.000000, -0.000000],
 [-0.000000, 0.000000],
 [0.082165, -0.246494],
 [0.000000, -0.000000],
 [0.000000, -0.000000]]

The first printed value is the squared error of the initial network at one input; it is not the mean training loss from the earlier log. The matrix that follows has exactly the shape of W1, so each entry answers how changing that weight locally changes this scalar loss. The upstream cotangent 2 * (y - t) already includes differentiation of the loss. Supplying one instead would ask for the derivative of the prediction. This distinction is easy to miss when calling a backward routine directly, because both cotangents have the same tensor shape.

Seven rows are zero because their ReLU units are inactive. Signed zeros reflect the arithmetic path; their signs alone would not distinguish exact zeros from underflowed values. Only row five survives, and its two entries stand in ratio -3, which is x_2/x_1=-0.75/0.25. If \delta_i denotes the loss derivative with respect to unit i's pre-activation, then \partial L/\partial (W_1)_{ij}=\delta_i x_j: each row is a scalar multiple of the input vector.

The claim to test is that 0.082165 is really a derivative. Perturb that one entry up and down and take a central difference of the loss:

some 0.082165

some (-0.246494)

Both entries agree with the backward pass to every printed digit. Each comparison changes one weight while holding the other weights, biases, input, and target fixed. Dividing by 2 * delta measures the change across the symmetric interval around that weight.

Tensor.modify? returns an Option because the coordinates are runtime naturals. The do block propagates a failed lookup, and some confirms that these coordinates existed. Step size needs a separate check: all calls here use nonzero delta, but the helper does not enforce that condition or decide whether the perturbation is small enough.

[0.012668, 0.004223, 0.080228]

[0.080228, 177.427282, 2035.348794]

some 0.000000

some 0.000000

some 0.002521

# Copy the Lean payload into double tensors before checking
# derivatives by finite differences.
# Python / PyTorch
import torch
torch.manual_seed(2026)

net = torch.nn.Sequential(
    torch.nn.Linear(2, 8),
    torch.nn.ReLU(),
    torch.nn.Linear(8, 1),
).double()

x = torch.tensor([0.25, -0.75], dtype=torch.float64)
t = torch.tensor([0.2], dtype=torch.float64)

def loss(w1):
    hidden = torch.relu(torch.nn.functional.linear(x, w1, net[0].bias))
    prediction = torch.nn.functional.linear(hidden, net[2].weight, net[2].bias)
    return ((prediction - t) ** 2).sum()

w1 = net[0].weight.detach().clone()
w1.requires_grad_(True)
print(torch.autograd.gradcheck(loss, (w1,), eps=1e-6))

Device Selection🔗

The source model does not mention CPU or CUDA. Device selection belongs to the runtime configuration:

# Run the training example with a CPU execution request.
lake exe torchlean quickstart_mlp --device cpu --steps 200

or, in a CUDA-enabled build:

# Build with CUDA support and request CUDA execution for the
# same example.
lake -R -K cuda=true exe torchlean \
  quickstart_mlp --device cuda --steps 200

The model type and parameter layout stay put. The selected backend profile plans the operations using the capsules available in that build, and --show-backend prints the resulting plan. Each line names the provider, the trust level, the VJP source, and how the operation is justified:

  matmul: reference.matmul provider=reference trust=checked vjp=torchlean-tape
    shape: shape safety for matmul; guarded at runtime by portable runtime shape checks
    layout: canonical-tensor layout compatibility for matmul

The complete report includes every capsule used by the run. These entries connect the device choice to the implementations of individual operations; the backend chapter explains how the planner selects them and what each trust level means.

Lowering The Forward Map🔗

Verification starts from an initialized model and a concrete parameter payload:

Verification.lowerForwardToIR reInit (nn.initialState reInit) : Except String (NN.Verification.Builtin.LoweredIR Float)

On success, LoweredIR Float pairs an NN.IR.Graph with its parameter payload store and the distinguished input and output nodes. The graph records shapes and operation tags; the payload's scalar type fixes how its entries are represented. The Except also admits a lowering error, which the examples below propagate through a do block.

The lowered graph expands the two linear layers and ReLU into primitive operations, parameter constants, and shape adjustments. Its node count is:

Except.ok 18

Eighteen nodes for a three-layer network, because lowering is explicit: parameters become constant nodes, a linear layer becomes a reshape, a matmul, and an add, and each node records its parents and output shape. The IR chapter walks the node list.

Using nn.initialState reInit analyzes the initial model. To analyze the trained model, the lowering must receive the trained runtime parameters through the lower-level manual interface or a saved exact-bits payload. Reusing the initial payload after training would verify another function. Everything below therefore describes the network at seed 2026, before any Adam step.

Input Regions🔗

A forward prediction answers "what did the model return at x?" Verification usually asks a quantified question. Around the held-out point

c=(0.25,-0.75),

an \ell_\infty ball of radius \varepsilon is the box

B_\varepsilon(c) =\{x\mid |x_i-c_i|\leq\varepsilon\text{ for }i=1,2\}.

FlatBox.lInfBall builds that box, LoweredIR.seedInputBox places it at the graph's input node, and runIBP propagates one interval per coordinate through the graph (Gowal et al., 2018). The output shape is checked before returning the bounds:

The if h : out.dim = 1 branch supplies evidence that the flat output box has one coordinate. The expressions h ▸ out.lo and h ▸ out.hi use it to return tensors of shape [1], preserving their entries without rounding them again. If the dimension differs, the function returns an error.

At radius zero, the input box contains only the held-out point, so we can compare its output bounds with the earlier prediction:

Except.ok ([-0.088261], [-0.088261])

Both endpoints print as -0.088261, matching rePred at this point and precision. This is a useful regression check of lowering and payload handling, not a proof that the functions agree on all inputs. Outward rounding can also give a nonzero-width enclosure at a single point.

Increasing the radius to 0.05 includes a neighborhood of that point:

Except.ok ([-0.111691], [-0.064831])

These are the computed candidate endpoints, rounded for display. Interpreting them as an enclosure for every point in the input box requires a soundness theorem for the chosen graph and arithmetic. The printed decimals themselves must not be substituted for outward-rounded exact endpoint bits.

[0.937198, 0.937198, 0.937198]

[0.730910, 1.276392]

Verification Parameters🔗

The training log ended with a held-out prediction of 0.228325, while the last interval calculation enclosed negative outputs. Both used the same architecture and input point. The difference is in the parameter payload: training updated the weights, but reBox still lowers nn.initialState reInit. The negative interval therefore describes the initial network.

To bound the trained network, we need its current parameter values, then we must repeat lowering and bound propagation with that payload. A saved verification result must identify those values along with its input region and arithmetic. The tensor shapes let the forward program accept the new payload; they cannot tell us whether it is the payload a particular claim was proved about.