4.2. GraphSpec: One Architecture, Several Meanings
The pure MLP formula from the previous chapter is excellent for a theorem:
x\mapsto W_2\operatorname{ReLU}(W_1x+b_1)+b_2.
It is less convenient for a tool that wants to enumerate layers, export parameter shapes, replace one operation, or lower the same architecture to several targets. For those tasks TorchLean uses GraphSpec, a small typed language in which the architecture itself is data.
GraphSpec sits between two other layers:
application model builders
↓
GraphSpec architecture and parameter ABI
↓
pure interpretation / TorchLean program / sequential model / DAG tools
↓
canonical NN.IR.Graph and runtime-specific lowering
GraphSpec preserves model structure and makes every input, output, and parameter shape explicit. The tensor runtime sits above it, while the low-level backend IR sits below it.
4.2.1. Write The Running MLP As A Graph
The complete architecture is:
import NN.GraphSpec.Models.Mlp
open NN
open NN.GraphSpec
open NN.GraphSpec.Models
open Spec
def mlpGraph (input hidden output : Nat) :
Graph
[ shape![hidden, input], shape![hidden],
shape![output, hidden], shape![output] ]
(shape![input])
(shape![output]) :=
Graph.linear input hidden >>>
Graph.relu (shape![hidden]) >>>
Graph.linear hidden output
Read the type from right to left:
-
the graph consumes one tensor of shape
[input]; -
it produces one tensor of shape
[output]; -
its parameter environment contains exactly four tensors;
-
the order is
W_1,b_1,W_2,b_2.
For the 2\to3\to1 model used by the GraphSpec tutorial, the parameter shapes are
W₁ : [3, 2] six values b₁ : [3] three values W₂ : [1, 3] three values b₂ : [1] one value
There are thirteen trainable scalars. That count is not recovered from strings such as
"layer1.weight". It follows from the graph's type.
4.2.2. Composition Computes The ABI
The composition operator >>> does more than connect two arrows. If
g₁ : Graph ps₁ σ τ g₂ : Graph ps₂ τ υ
then
g₁ >>> g₂ : Graph (ps₁ ++ ps₂) σ υ.
The intermediate shape must be the same \tau, and the parameter lists are concatenated in
construction order. Replacing the first linear layer by Graph.linear input 5 changes its output
shape to [5]; the existing ReLU can still consume it, but the second linear layer must accept
five inputs. Lean reports the mismatch at the architecture definition.
This is the first hands-on experiment:
-
change the hidden dimension at only one of the two linear nodes;
-
ask Lean to elaborate the file.
The failure occurs before initialization, data loading, or runtime execution because the broken edge has no well-typed composition.
4.2.3. A Primitive Has Two Interpretations
A sequential GraphSpec primitive stores:
specFwd : pure shape-indexed tensor function torchProgram : executable TorchLean frontend program
For a linear primitive, specFwd calls the mathematical linearSpec; torchProgram constructs the
runtime operation with the same input and parameter shapes. Composition builds both interpretations
in parallel.
The record does not prove that the two fields agree merely by storing them together. The important difference is:
-
construction keeps the intended programs adjacent;
-
a theorem establishes that their meanings coincide.
TorchLean proves such relationships at the primitive or model level where they are available. This avoids a global axiom saying that every future GraphSpec operation is correct by construction.
4.2.4. Current Primitive Adapters
The minimal sequential vocabulary in NN.GraphSpec.Core is linear, relu, and softmax.
NN.GraphSpec.Primitives.Vision adds CHW conv2d, maxPool2d, flatten, and channel-first
BatchNorm. Each of these adapters supplies the same three pieces of information:
-
the exact parameter-shape list;
-
a pure
specFwdmeaning; -
an executable
torchProgrammeaning.
The MLP and two-convolution CNN compose the sequential adapters. residualLinear demonstrates the
DAG language and an explicit skip connection. These are implemented examples, not a claim that
every layer under NN.Spec already has a GraphSpec adapter. GraphSpec.ToTorchLean.toSeq is also
deliberately partial: it succeeds only when every primitive provides a corresponding layer
constructor.
4.2.5. Run The Complete Lowering
The repository contains an executable GraphSpec tutorial:
lake exe torchlean graphspec --device cpu --backend eager
The current checkout prints:
== GraphSpec tutorial == GraphSpec architecture ladder: 1. MLP: sequential layer stack; lowers to nn.Sequential and trains below. 2. CNN: sequential vision graph with checked conv/pool shape arithmetic. 3. residualLinear: minimal DAG-native skip connection. model: Sequential: [2] -> [1], layers=3, params=13 [0] Linear(2, 3): [2] -> [3] params=9 [[3, 2], [3]] [1] ReLU: [3] -> [3] params=0 [] [2] Linear(3, 1): [3] -> [1] params=4 [[1, 3], [1]] dataset size = 1 mean_loss(before) = 1.239197 mean_loss(after) = 0.247518 steps=3 loss0=1.239197 loss1=0.247518
The execution path is:
Models.mlp │ Graph ps [2] [1] ▼ GraphSpec.ToTorchLean.toSeq │ Except String (nn.Sequential [2] [1]) ▼ Trainer.new ▼ eager runtime and autograd tape
toSeq is intentionally partial. A sequential linear/ReLU stack has an nn layer counterpart, so
the conversion succeeds. An arbitrary custom primitive may have a pure and program interpretation
without having an nn.LayerDef constructor; in that case the conversion returns an error rather
than inventing a layer.
Try the compiled frontend as a second run:
lake exe torchlean graphspec --device cpu --backend compiled
The architecture and parameter ABI are unchanged. Only the execution path selected after lowering changes.
4.2.6. Pure Interpretation
GraphSpec can be evaluated without the trainer:
Interp.spec mlpGraph params x
Here params is a heterogeneous tensor list whose shape index is exactly
[[hidden, input], [hidden], [output, hidden], [output]].
Pattern matching on that list reveals W_1, b_1, W_2, and b_2 in ABI order. The
interpreter then computes
\operatorname{linearSpec}
(W_2,b_2)
\left(\operatorname{ReLU}
\left(\operatorname{linearSpec}(W_1,b_1,x)\right)\right).
The theorem
NN.GraphSpec.Models.mlp_interp_eq_spec_mlp_forward proves that this interpretation equals
TorchLean's hand-written MLP specification for every scalar \alpha satisfying the required Context,
every well-shaped parameter list, and every input.
That is stronger than checking a few Float examples: it is a universally quantified equality of the two pure definitions. It is also narrower than runtime correctness: neither CUDA buffers nor an eager tape occur in the theorem.
4.2.7. Why A Second, DAG-Shaped Syntax Exists
A chain can be written with >>>. A residual block cannot:
r(x)=\operatorname{ReLU}(Wx+b+x).
The input x is used twice. Hiding this in a special ResidualLinear primitive would make one
example work, but every new sharing pattern would demand another special primitive. GraphSpec's DAG
language instead represents sharing directly.
A term has type
DAG.Term Γ τ
meaning that, given a typed environment \Gamma, it computes a tensor of shape \tau. Its essential
constructors are:
-
var, which reads an existing value; -
op, which applies an arbitrary-arity primitive; -
let1, which computes an intermediate once and extends the environment.
The residual computation is conceptually:
let y = linear(W, b, x) let z = add(y, x) relu(z)
The same x variable appears in the linear and add arguments, while y is bound once. Because the
language has no recursion and only extends the environment with prior values, its terms denote
acyclic graphs by construction.
Run the tutorial command again and notice the third architecture in its printed ladder:
residualLinear is checked as a DAG model even though the short tutorial trains only the sequential
MLP.
4.2.8. Sequential Graph Versus DAG Model
The two GraphSpec syntaxes are related but not identical:
Form | Best use | Sharing | Parameter representation |
|---|---|---|---|
| readable layer chains | no explicit fan-out |
type-indexed list |
| residual and multi-input models |
explicit variables and | typed model environment |
A sequential graph can be lowered structurally into a DAG model. This does not require a numerical theorem because the conversion also comes with a theorem relating the relevant pure interpretation. The reverse direction is not generally possible: a DAG with fan-out has no faithful representation as a plain chain without adding duplication or a special combinator.
4.2.9. Initialization Is Separate From Meaning
GraphSpec knows the parameter shapes and order, but an architecture does not mathematically require one initializer. The repository supports zero initialization for simple structural examples and deterministic seeded initialization for executable models.
For the MLP, the theorem
mlp_detInitParams_eq_torchlean_linear_inits checks that GraphSpec's deterministic traversal
produces parameters in the same ABI order as the corresponding TorchLean linear-layer
initializers. It does not say that those random-looking values are optimal, and it says nothing
about an external checkpoint.
An imported checkpoint needs a different argument:
-
parse the external names and arrays;
-
check each concrete shape and finite-value condition;
-
map values into the GraphSpec ABI;
-
state or assume that this mapping matches the source framework's layout.
Architecture correctness cannot authenticate training provenance.
4.2.10. GraphSpec Is Not The Backend IR
It is tempting to call every graph representation “the graph,” but GraphSpec and NN.IR.Graph
have different jobs.
GraphSpec is intrinsically typed. A model that connects incompatible shapes does not elaborate. Its primitives carry pure and TorchLean-program interpretations.
NN.IR.Graph is a serializable op-tagged DAG. Nodes carry numeric IDs, parent IDs, output shapes,
and attributes; tensors and parameters live in an external payload. This form is better for
importers, validators, generic passes, verification, and backend planning.
There is not currently a universal compiler from every GraphSpec model to canonical IR. Selected frontend and model paths lower to IR, and selected semantic theorems cover those paths. The next chapter builds and inspects that lower-level representation directly.