GraphSpec model: CNN (2 convs)

This is a GraphSpec version of the classic small CNN:

Conv2D → ReLU → MaxPool2D → Conv2D → ReLU → MaxPool2D → Flatten → Linear

Notes:

• This is still a chain model (no skip connections), so it fits the sequential GraphSpec core (NN.GraphSpec.Core) and can be built using >>>.
• Parameter shapes are derived in the type: (K1,b1,K2,b2,W,b) where K are Conv kernels and W is the linear head matrix.

Shape bookkeeping (why all the Nat arithmetic?)

The point of GraphSpec is that the shape interface is part of the type. Convolution and pooling therefore bake their output shapes into the type, using the standard formulas:

• Conv2D: outH = Spec.Shape.slidingWindowOutDim inH kH stride padding (and similarly for outW)
• MaxPool2D: outH = Spec.poolOutDim inH kH stride 0 (and similarly for outW)

This file defines small helper abbreviations (outH/outW/poolH/poolW) so that the overall classifier head shape (the input dimension to the final Linear) is computed once and reused.

Why not put this under GraphSpec.DAG?

You can express a chain model as a DAG term, but the sequential DSL is the simpler interface for pure pipelines. The DAG language is reserved for models that need sharing or multi-input nodes (e.g. residual adds).

Conv/pool output size helpers

@[reducible, inline]

abbrev NN.GraphSpec.Models.outH (inH kH stride padding : ℕ) : ℕ

Convolution output height formula (standard DL convention).

@[reducible, inline]

abbrev NN.GraphSpec.Models.outW (inW kW stride padding : ℕ) : ℕ

Convolution output width formula (standard DL convention).

@[reducible, inline]

abbrev NN.GraphSpec.Models.poolH (inH kH stride : ℕ) : ℕ

MaxPool output height formula.

@[reducible, inline]

abbrev NN.GraphSpec.Models.poolW (inW kW stride : ℕ) : ℕ

MaxPool output width formula.

@[reducible, inline]

abbrev NN.GraphSpec.Models.featH (inH kH stride1 padding1 poolKH poolStride1 stride2 padding2 poolStride2 : ℕ) : ℕ

Final feature-map height after:

conv1 → pool1 → conv2 → pool2.

@[reducible, inline]

abbrev NN.GraphSpec.Models.featW (inW kW stride1 padding1 poolKW poolStride1 stride2 padding2 poolStride2 : ℕ) : ℕ

Final feature-map width after:

conv1 → pool1 → conv2 → pool2.

@[reducible, inline]

abbrev NN.GraphSpec.Models.featSize (c2 inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2 : ℕ) : ℕ

Total flattened feature size for the classifier head.

If the second conv produces c2 channels and the final spatial size is featH × featW, then the flattened vector length is (CHW c2 featH featW).size.

def NN.GraphSpec.Models.twoConvCnn (inC c1 c2 outDim inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2 : ℕ) {h_inC : inC ≠ 0} {h_c1 : c1 ≠ 0} {_h_c2 : c2 ≠ 0} {h_kH : kH ≠ 0} {h_kW : kW ≠ 0} {h_stride1 : stride1 ≠ 0} {h_stride2 : stride2 ≠ 0} {h_poolKH : poolKH ≠ 0} {h_poolKW : poolKW ≠ 0} {h_poolStride1 : poolStride1 ≠ 0} {h_poolStride2 : poolStride2 ≠ 0} : Graph [Spec.Shape.dim c1 (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))), Spec.Shape.dim c1 Spec.Shape.scalar, Spec.Shape.dim c2 (Spec.Shape.dim c1 (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))), Spec.Shape.dim c2 Spec.Shape.scalar, Spec.Shape.dim outDim (Spec.Shape.dim (featSize c2 inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2) Spec.Shape.scalar), Spec.Shape.dim outDim Spec.Shape.scalar] (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) (Spec.Shape.dim outDim Spec.Shape.scalar)

2-conv CNN GraphSpec model.

Conventions:

• both conv layers share (kH,kW) but can differ in stride/padding.
• pooling uses kernel (poolKH,poolKW) and can differ in stride per pooling site.

Parameter layout (type-level):

• (K1,b1) for conv1,
• (K2,b2) for conv2,
• (W,b) for the final linear head.

Input/output:

• input x : CHW inC inH inW (no batch dimension),
• output logits : Vec outDim.

def NN.GraphSpec.Models.twoConvCnnDAGModelZeroInit (inC c1 c2 outDim inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2 : ℕ) {h_inC : inC ≠ 0} {h_c1 : c1 ≠ 0} {h_c2 : c2 ≠ 0} {h_kH : kH ≠ 0} {h_kW : kW ≠ 0} {h_stride1 : stride1 ≠ 0} {h_stride2 : stride2 ≠ 0} {h_poolKH : poolKH ≠ 0} {h_poolKW : poolKW ≠ 0} {h_poolStride1 : poolStride1 ≠ 0} {h_poolStride2 : poolStride2 ≠ 0} : DAG.Model [Spec.Shape.dim c1 (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))), Spec.Shape.dim c1 Spec.Shape.scalar, Spec.Shape.dim c2 (Spec.Shape.dim c1 (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))), Spec.Shape.dim c2 Spec.Shape.scalar, Spec.Shape.dim outDim (Spec.Shape.dim (featSize c2 inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2) Spec.Shape.scalar), Spec.Shape.dim outDim Spec.Shape.scalar] [Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))] (Spec.Shape.dim outDim Spec.Shape.scalar)

The same 2-conv CNN, but exposed as a DAG Model via the structural lowering LowerToDAG.Graph.toDAGModelZeroInit.

This lets “DAG-only” downstream tooling consume this architecture even though it is authored as a sequential Graph pipeline.

Initialization: all-zero parameters (see LowerToDAG.Graph.toDAGModelZeroInit).