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:

nn.Sequential!: layer 2 expects input shape [7], but layer 1 outputs [8].
Change layer 2's input shape or insert a layer that converts [8] to [7].

The final layer constrains ReLU's shape to [7], so the diagnostic points to layer 2: that ReLU cannot consume the first layer's [8] output. 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 Float : 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 := mse
      optimizer := optim.adam { lr := 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.
scripts/lake.sh exe torchlean quickstart_mlp \
  --device cpu \
  --steps 200 \
  --seed 2026

The command reports the dataset size, mean loss before and after training, periodic step losses, and a prediction at the held-out input. Read the two loss reports separately: mean_loss evaluates all 25 examples, while a step loss measures the single example used for that update. A low step loss can therefore coexist with a much larger dataset mean.

Compare trained(heldout) with the target 0.2 as well. That prediction tests one point between grid locations; it answers a different question from the average error on the training grid.

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

scripts/lake.sh exe torchlean quickstart_mlp --device cpu --steps 200 \
  --seed 2026 --arithmetic ieee

python3 - <<'PY'
# Compare a full-batch PyTorch run; its initial weights and
# update schedule differ.
import torch
from torch import nn

torch.manual_seed(2026)
model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 1))
axis = torch.linspace(-1, 1, 5)
grid = torch.cartesian_prod(axis, axis)
ys = (0.8 * torch.relu(grid[:, 0] + grid[:, 1])
      - 0.4 * torch.relu(grid[:, 1] - grid[:, 0]) + 0.2).unsqueeze(1)
heldout = torch.tensor([0.25, -0.75])
opt = torch.optim.Adam(model.parameters(), lr=0.03)
with torch.no_grad():
    print("mean_loss(before) =", nn.MSELoss()(model(grid), ys).item())
for step in range(200):
    opt.zero_grad()
    loss = nn.MSELoss()(model(grid), ys)
    loss.backward()
    opt.step()
with torch.no_grad():
    print("mean_loss(after) =", nn.MSELoss()(model(grid), ys).item())
    print("trained(heldout) =", model(heldout).item())
PY

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. Compare this result with untrained(heldout) from the command above: both start from the seed-2026 payload, but this calculation uses host Float while the trainer selects its configured arithmetic. If they disagree, this small forward calculation helps separate an error in the MLP equations from an error in loading or executing the state, or a difference in rounding. 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 target does not require all eight hidden units. 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. If a training run leaves a nonzero error, insufficient width alone cannot explain it: the architecture can contain this two-unit construction. 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 reported by the command. 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

# Initialize a separate PyTorch model in double precision
# and check its 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.
scripts/lake.sh 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.
scripts/lake.sh -R -Kcuda=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 lists the capsules selected by the profile. It describes the intended implementations, not evidence that each native kernel executed. 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 Float) : 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. For a supported training result, call trained.verify to analyze its snapshot, or read trained.state when constructing a graph-level workflow. A saved exact-bits checkpoint is another way to retain that payload. Reusing the initial state after training would analyze 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 interval calculation encloses negative outputs because reBox lowers nn.initialState reInit. Training updates a separate parameter payload. A changed held-out prediction from the trainer therefore cannot be compared with this box as though the two described the same fitted function. The architecture and input point agree; the weights need not.

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.