TorchLean

4.1. The Mathematical Specification🔗

The model we trained earlier can be written on one line:

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

That formula is the mathematical center of the example. It says nothing about mutable buffers, CUDA launches, PyTorch modules, reverse-mode tapes, or JSON files. It does say exactly how the four parameter tensors and the input determine the output. TorchLean's NN.Spec library is where such formulas live as Lean definitions.

The distinction is important because executable ML code changes form constantly. A linear layer may be evaluated by a nested Lean function, a CPU loop, cuBLAS, or an ATen kernel. The specification gives all of those implementations one statement to satisfy.

4.1.1. The Tensor Behind The Formula🔗

A specification tensor is defined recursively from its shape:

  • a scalar shape stores one value;

  • a dimension of length n stores a function from Fin n to a smaller tensor.

Thus a vector of length two is, mathematically, two scalar values indexed by Fin 2; a matrix of shape [3, 2] is three such vectors. The index type prevents an out-of-range lookup. It also makes shape induction natural: a proof about an arbitrary tensor can follow the same scalar-or-dimension recursion as the datatype.

This is not a claim that CUDA stores a matrix as nested Lean functions. Native runtimes use flat, contiguous buffers. The specification chooses the representation that makes mathematical reasoning clear; a layout contract is needed when an implementation flattens that value into memory.

The exported aliases hide most of the recursive spelling:

import NN.Spec.Layers.Linear
import NN.Spec.Layers.Activation

open Spec
open Spec.Tensor

#check Tensor ℝ (shape![2])
#check Tensor ℝ (shape![3, 2])
#check LinearSpec ℝ 2 3

For a linear map from two inputs to three outputs, LinearSpec ℝ 2 3 contains

W\in\mathbb{R}^{3\times 2}, \qquad b\in\mathbb{R}^{3}.

The order is the same convention used by PyTorch's nn.Linear: output features index the rows of the weight matrix.

4.1.2. One Layer, Forward And Backward🔗

The forward definition is deliberately short:

def linearSpec {α : Type} [Add α] [Mul α] [Zero α]
    {inDim outDim : Nat}
    (layer : LinearSpec α inDim outDim)
    (x : Tensor α (shape![inDim])) :
    Tensor α (shape![outDim]) :=
  addSpec (matVecMulSpec layer.weights x) layer.bias

At coordinate i, this means

y_i=b_i+\sum_{j=0}^{\mathrm{inDim}-1}W_{ij}x_j.

The backward specification receives an upstream cotangent g=\partial L/\partial y and returns

\frac{\partial L}{\partial W}=g\,x^\mathsf{T},\qquad \frac{\partial L}{\partial b}=g,\qquad \frac{\partial L}{\partial x}=W^\mathsf{T}g.

In Lean these three tensors have shapes [outDim, inDim], [outDim], and [inDim]. Returning a transposed weight gradient by accident is therefore a type error rather than a numerically plausible array.

The source definitions are linearSpec and linearBackwardSpec in the same file. Later, the autograd proofs compare executable VJP rules with these definitions. The existence of the definitions alone does not prove the comparison.

4.1.3. The Complete MLP Meaning🔗

The running model is the composition of two linear specifications and a pointwise activation:

def twoLayerMlp
    {α : Type} [Context α]
    {inDim hidden outDim : Nat}
    (first : LinearSpec α inDim hidden)
    (second : LinearSpec α hidden outDim)
    (x : Tensor α (shape![inDim])) :
    Tensor α (shape![outDim]) :=
  linearSpec second
    (Activation.reluSpec (linearSpec first x))

Expanding this definition gives the formula at the start of the chapter. No graph traversal or runtime state is hidden inside it.

The ReLU definition commits to two choices:

\operatorname{ReLU}(z)=\max(z,0),

and for the selected derivative,

\operatorname{ReLU}'(z)= \begin{cases} 1 & z>0,\\ 0 & z\le 0. \end{cases}

The value at the kink matters. Other subgradients are mathematically defensible, but a forward and backward correctness theorem needs one concrete rule. TorchLean chooses zero at the kink, matching the rule used by the current runtime path.

4.1.4. Run The Same Formula🔗

The quickest executable using this architecture is:

lake exe torchlean quickstart_mlp --device cpu --steps 200 --seed 2026

On the current example dataset it reports:

dataset size = 25
mean_loss(before) = 0.761530
mean_loss(after) = 0.003234
heldout x=(0.25,-0.75), target=0.2, prediction(after)=[0.210239]

The command executes runtime tensors and an autograd tape; it does not evaluate twoLayerMlp by reducing the pure Lean definition above. The connection is made operation by operation:

  1. the nn.linear builder fixes the same parameter shapes;

  2. its forward program emits the runtime linear operation;

  3. the graph interpreter assigns .linear the linearSpec denotation;

  4. the VJP proof identifies the selected backward rule with linearBackwardSpec;

  5. a backend capsule records which native provider, if any, executed the operation.

This chain is why a specification is useful. It gives each bridge a stable target.

4.1.5. Change One Value🔗

The formula can be inspected without training. Take

W_1=\begin{bmatrix}1&1\\-1&1\end{bmatrix}, \quad b_1=0,\quad W_2=\begin{bmatrix}0.8&-0.4\end{bmatrix}, \quad b_2=0.2.

For x=(0.25,-0.75),

W_1x=(-0.5,-1),\qquad \operatorname{ReLU}(W_1x)=(0,0),

so the exact-real output is 0.2. Change only the first bias to 0.6. The first hidden preactivation becomes 0.1, and the output becomes

0.8(0.1)+0.2=0.28.

This tiny calculation is the same kind of reasoning used in interval propagation: replace one point by a set of possible inputs, then bound every intermediate tensor.

4.1.6. Scalar Polymorphism Is A Real Choice🔗

The type parameter \alpha determines what the symbols +, *, max, exp, and division mean. TorchLean reuses the tensor structure at several scalar interpretations:

Scalar

Meaning

exact real arithmetic used for mathematical statements

FP32

binary32-precision, gradual-underflow rounded-real values with no upper exponent cutoff

IEEE32Exec

executable binary32 bit patterns, including signed zero, infinity, and NaN

interval contexts

sets of possible values, with outward enclosure operations

runtime Float

Lean's native executable floating-point value

Writing one polymorphic definition is not a proof that these interpretations agree. For example,

\operatorname{softmax}(x)_i= \frac{\exp(x_i-m)}{\sum_j\exp(x_j-m)},\qquad m=\max_jx_j

is a clean real-valued formula. A binary32 implementation introduces rounding in max, subtraction, exponential approximation, summation, and division. A CUDA reduction may also choose a different summation tree. The later floating-point and runtime-approximation chapters state the conditions under which one interpretation encloses or approximates another.

4.1.7. The Implemented Specifications🔗

The spec layer is broader than the running MLP. NN.Spec.Layers contains typed meanings for linear algebra, convolution and transposed convolution, two-dimensional and N-dimensional pooling, activations and losses, normalization, dropout, embeddings, recurrent cells, selective scan, and scaled dot-product attention. Model definitions under NN.Spec.Models compose these operations into families such as CNNs, transformers, recurrent networks, and state-space models.

That inventory is a semantics inventory, not a runtime support matrix. An operation can have a pure definition before it has an eager tape rule, a compiled lowering, a CUDA kernel, or an end-to-end correctness theorem. The relevant runtime and compiler chapters name those smaller supported fragments explicitly.

4.1.8. Shape Indices And Scalar Types🔗

The shape index s and scalar parameter \alpha answer different questions. Tensor α s fixes both for one value, and Graph.denote follows the homogeneous-scalar contract from Tensors And Shapes. The later IR chapter uses “heterogeneous” only for a table of differently shaped values, not for a mixed-dtype graph.

4.1.9. Conventions That Must Be In The Definition🔗

Shapes are only one source of ambiguity. The spec layer also fixes choices that a model name does not determine.

4.1.9.1. Loss reductions🔗

mseSpec takes a global mean over all entries. Cross-entropy over logits applies log-softmax on the final class axis and averages over the remaining slices. Changing mean to sum changes both the loss and every gradient by a scale factor.

4.1.9.2. Attention masks🔗

A boolean attention mask is a hard support constraint:

  • true allows the key;

  • false gives the key exactly zero softmax numerator;

  • if every key in a row is blocked, the output row is zero.

This is the finite formulation of a negative-infinity mask. It is not an additive score bias of -1000 or any other finite sentinel. For sufficiently large logits a finite sentinel can leak nonzero probability into a blocked position; the hard-mask definition cannot.

4.1.9.3. Dropout🔗

Randomness is explicit. A masked dropout specification receives the mask as an argument. Runtime training code may generate that mask from a seed and tape state, but the semantic function does not consult hidden global randomness.

4.1.9.4. Invalid windows🔗

Some mathematical operations are total where the runtime is partial. A convolution output extent is normally

\left\lfloor\frac{n+2p-k}{s}\right\rfloor+1.

The shape helper defines an answer even around invalid kernel or stride configurations, while the runtime validator rejects unsupported calls before reaching native code. A total denotation and an admission check answer different questions.

4.1.10. A Proof Checkpoint🔗

GraphSpec's checked MLP uses the same four parameter tensors:

[W₁ : shape![hidden, input],
 b₁ : shape![hidden],
 W₂ : shape![output, hidden],
 b₂ : shape![output]]

The theorem NN.GraphSpec.Models.mlp_interp_eq_spec_mlp_forward proves that interpreting that GraphSpec model is equal to the hand-written two-layer specification. Its conclusion is about two pure meanings. It does not mention the eager tape, CUDA, or an imported checkpoint, so it should not be reported as a proof of those objects.

That theorem supplies the first exact link in a longer chain and prevents later engineering layers from silently changing what “the MLP” means.

The next chapter turns these formulas into typed architectures with explicit parameter layouts.