Spec helpers: param counting, feature queries, arch display, validation.

def fpnDetectParamShapes (oc c3 c4 c5 A tower : Nat) : List (List Nat)

Param shapes owned by a .fpnDetect layer, in the ONE canonical order that emitTrainStepSig, the optimizer, the grad-clip list, the backward's grad-binding, NetSpec.paramShapes and heInitLayer must all agree on:

Wn3 [oc,c3], Wn4 [oc,c4], Wn5 [oc,c5]           -- neck laterals
(W [oc,oc,3,3], b [oc]) × tower, per level P3 → P4 → P5
Wh3, Wh4, Wh5 [A·15, oc]                        -- head convs
bh3, bh4, bh5 [A·15]                            -- head biases, LAST

The biases MUST stay last: NetSpec.applyDetPriorBias installs the RetinaNet prior by splicing the final 3·A·15 floats of the buffer.

Defined here (rather than in the codegen) because SpecHelpers and MlirCodegen both need it and neither should own it — a mismatch between any two of these consumers is a silent parameter-aliasing bug, not a compile error.

Equations

• One or more equations did not get rendered due to their size.

def Layer.nParams : Layer → Nat

Equations

• One or more equations did not get rendered due to their size.
• (Layer.conv2d ic oc k pad act).nParams = oc * ic * k * k + oc
• (Layer.convBn ic oc k stride pad).nParams = oc * ic * k * k + 2 * oc
• (Layer.layerNorm d).nParams = 2 * d
• (Layer.dense fi fo act).nParams = fi * fo + fo
• (Layer.separableConv ic oc stride).nParams = ic * 9 + 2 * ic + (oc * ic + 2 * oc)
• (Layer.fireModule ic sq e1 e3).nParams = sq * ic + 2 * sq + (e1 * sq + 2 * e1) + (e3 * sq * 9 + 2 * e3)
• (Layer.patchEmbed ic dim p nP).nParams = dim * ic * p * p + dim + dim + (nP + 1) * dim
• (Layer.mambaBlock dim stateSize expand nBlocks).nParams = nBlocks * (3 * expand * dim * dim + 3 * expand * dim * stateSize + 5 * expand * dim + dim)
• (Layer.patchMerging inDim outDim).nParams = 2 * 4 * inDim + (4 * inDim * outDim + outDim)
• (Layer.unetDown ic oc).nParams = 9 * ic * oc + 2 * oc + (9 * oc * oc + 2 * oc)
• (Layer.unetUp ic oc).nParams = 9 * (ic + oc) * oc + 2 * oc + (9 * oc * oc + 2 * oc)
• (Layer.transformerDecoder dim _heads mlpDim nBlocks nQueries).nParams = nBlocks * (3 * 2 * dim + 2 * (4 * (dim * dim + dim)) + (dim * mlpDim + mlpDim) + (mlpDim * dim + dim)) + nQueries * dim
• (Layer.detrHeads dim nClasses).nParams = dim * (nClasses + 1) + (nClasses + 1) + (dim * dim + dim) + (dim * dim + dim) + (dim * 4 + 4)
• (Layer.shuffleV2Block ic oc nUnits).nParams = 12 * ic + ic * oc + 11 * oc + oc * oc / 4 + (nUnits - 1) * (oc * oc / 2 + 19 * oc / 2)
• (Layer.asppModule ic oc).nParams = ic * oc + oc + 2 * oc + 3 * (9 * ic * oc + oc + 2 * oc) + (ic * oc + oc + 2 * oc) + (5 * oc * oc + oc + 2 * oc)
• (Layer.fpnModule c2 c3 c4 c5 target).nParams = (c2 + c3 + c4 + c5) * target + 4 * target + 4 * (9 * target * target + target)
• (Layer.evoformerBlock msaChannels pairChannels nBlocks).nParams = nBlocks * (16 * msaChannels * msaChannels + msaChannels * msaChannels * pairChannels / 32 + 16 * pairChannels * pairChannels)
• (Layer.structureModule singleChannels pairChannels nBlocks).nParams = 6 * singleChannels * singleChannels + singleChannels * pairChannels / 4
• (Layer.convNextStage channels nBlocks norm act).nParams = nBlocks * (8 * channels * channels + 58 * channels)
• (Layer.convNextDownsample ic oc norm).nParams = 2 * ic + 4 * ic * oc + oc
• (Layer.convNextStem ic oc p).nParams = oc * ic * p * p + oc + 2 * oc
• (Layer.waveNetBlock residualCh skipCh nLayers).nParams = nLayers * (4 * residualCh * residualCh + 2 * residualCh + residualCh * residualCh + residualCh + residualCh * skipCh + skipCh)
• (Layer.positionalEncoding _inputDim _numFrequencies).nParams = 0
• (Layer.darknetBlock channels nBlocks).nParams = nBlocks * (5 * channels * channels + 3 * channels)
• (Layer.inceptionModule ic b1 b2r b2 b3r b3 b4).nParams = ic * b1 + 2 * b1 + (ic * b2r + 2 * b2r + (9 * b2r * b2 + 2 * b2)) + (ic * b3r + 2 * b3r + (25 * b3r * b3 + 2 * b3)) + (ic * b4 + 2 * b4)
• (Layer.transitionLayer ic oc).nParams = 2 * ic + ic * oc + oc
• (Layer.tokenPositionEmbed v t d idsInput gather posEmb).nParams = if posEmb = true then v * d + t * d else v * d
• (Layer.lmHead d v seqLen).nParams = d * v + v
• (Layer.timeCondAdd c nFreq).nParams = 2 * nFreq * c + c
• x✝.nParams = 0

def NetSpec.totalParams (s : NetSpec) : Nat

Equations

• s.totalParams = List.foldl (fun (acc : Nat) (l : Layer) => acc + l.nParams) 0 s.layers

def NetSpec.hasConv (s : NetSpec) : Bool

Equations

• One or more equations did not get rendered due to their size.

def NetSpec.hasPool (s : NetSpec) : Bool

Equations

• s.hasPool = s.layers.any fun (x : Layer) => match x with | Layer.maxPool size stride => true | x => false

def NetSpec.hasBn (s : NetSpec) : Bool

Equations

• One or more equations did not get rendered due to their size.

def NetSpec.hasResidual (s : NetSpec) : Bool

Equations

• One or more equations did not get rendered due to their size.

def NetSpec.hasBottleneck (s : NetSpec) : Bool

Equations

• s.hasBottleneck = s.layers.any fun (x : Layer) => match x with | Layer.bottleneckBlock ic oc nBlocks firstStride => true | x => false

def NetSpec.hasSeparable (s : NetSpec) : Bool

Equations

• One or more equations did not get rendered due to their size.

def NetSpec.hasInvertedResidual (s : NetSpec) : Bool

Equations

• s.hasInvertedResidual = s.layers.any fun (x : Layer) => match x with | Layer.invertedResidual ic oc expand stride nBlocks => true | x => false

def NetSpec.hasMbConv (s : NetSpec) : Bool

Equations

• One or more equations did not get rendered due to their size.

def NetSpec.hasGlobalAvgPool (s : NetSpec) : Bool

Equations

• s.hasGlobalAvgPool = s.layers.any fun (x : Layer) => match x with | Layer.globalAvgPool => true | x => false

def NetSpec.hasTransformer (s : NetSpec) : Bool

Equations

• s.hasTransformer = s.layers.any fun (x : Layer) => match x with | Layer.transformerEncoder dim heads mlpDim nBlocks causalMask keepSequence flashAttn rope => true | x => false

def NetSpec.numClasses (s : NetSpec) : Nat

Equations

• s.numClasses = match s.layers.getLast? with | some (Layer.dense fanIn fo act) => fo | some (Layer.lmHead dModel v seqLen) => v | x => 0

def NetSpec.archStr (s : NetSpec) : String

Equations

• One or more equations did not get rendered due to their size.

def Layer.outChannels : Layer → Nat

Output channels of a layer. Returns 0 for structural layers (pool, flatten, GAP).

Equations

• (Layer.conv2d ic oc kSize pad act).outChannels = oc
• (Layer.convBn ic oc kSize stride pad).outChannels = oc
• (Layer.layerNorm d).outChannels = d
• (Layer.dense fanIn fo act).outChannels = fo
• (Layer.residualBlock ic oc nBlocks firstStride).outChannels = oc
• (Layer.bottleneckBlock ic oc nBlocks firstStride).outChannels = oc
• (Layer.separableConv ic oc stride).outChannels = oc
• (Layer.invertedResidual ic oc expand stride nBlocks).outChannels = oc
• (Layer.mbConv ic oc expand kSize stride nBlocks useSE act).outChannels = oc
• (Layer.mbConvV3 ic oc expandCh kSize stride useSE act).outChannels = oc
• (Layer.fusedMbConv ic oc expand kSize stride nBlocks useSE).outChannels = oc
• (Layer.uib ic oc expand stride preDWk postDWk).outChannels = oc
• (Layer.fireModule ic squeeze e1 e3).outChannels = e1 + e3
• (Layer.patchEmbed ic dim patchSize nPatches).outChannels = dim
• (Layer.transformerEncoder dim heads mlpDim nBlocks causalMask keepSequence flashAttn rope).outChannels = dim
• (Layer.mambaBlock dim stateSize expand nBlocks).outChannels = dim
• (Layer.swinStage dim heads mlpDim windowSize nBlocks).outChannels = dim
• (Layer.patchMerging inDim outDim).outChannels = outDim
• (Layer.unetDown ic oc).outChannels = oc
• (Layer.unetUp ic oc).outChannels = oc
• (Layer.transformerDecoder dim heads mlpDim nBlocks nQueries).outChannels = dim
• (Layer.detrHeads dim nClasses).outChannels = nClasses + 1
• (Layer.shuffleBlock ic oc groups nUnits).outChannels = oc
• (Layer.shuffleV2Block ic oc nUnits).outChannels = oc
• (Layer.asppModule ic oc).outChannels = oc
• (Layer.fpnModule c2 c3 c4 c5 target).outChannels = target
• (Layer.fpnDetect oc c3 c4 c5 g5 A tower).outChannels = 315 * A * g5 * g5
• (Layer.evoformerBlock msaCh pairChannels nBlocks).outChannels = msaCh
• (Layer.structureModule sCh pairChannels nBlocks).outChannels = sCh
• (Layer.mobileVitBlock ic dim heads mlpDim nTxBlocks).outChannels = ic
• (Layer.convNextStage c nBlocks norm act).outChannels = c
• (Layer.convNextDownsample ic oc norm).outChannels = oc
• (Layer.convNextStem ic oc patch).outChannels = oc
• (Layer.waveNetBlock residualCh skipCh nLayers).outChannels = skipCh
• (Layer.positionalEncoding inputDim numFreq).outChannels = inputDim * 2 * numFreq
• (Layer.nerfMLP encodedPosDim encodedDirDim hiddenDim).outChannels = 4
• (Layer.darknetBlock c nBlocks).outChannels = c
• (Layer.cspBlock ic oc nBlocks).outChannels = oc
• (Layer.inceptionModule ic b1 b2reduce b2 b3reduce b3 b4).outChannels = b1 + b2 + b3 + b4
• (Layer.denseBlock ic gr nLayers).outChannels = ic + nLayers * gr
• (Layer.transitionLayer ic oc).outChannels = oc
• (Layer.tokenPositionEmbed vocabSize seqLen d idsInput gather).outChannels = d
• (Layer.lmHead dModel v t).outChannels = t * v
• (Layer.timeCondAdd c nFreq).outChannels = c
• Layer.spatialFlatten.outChannels = 0
• (Layer.spatialUnflatten c height width).outChannels = c
• x✝.outChannels = 0

def Layer.inChannels : Layer → Nat

Input channels expected by a layer. Returns 0 for layers that accept any input.

Equations

• (Layer.conv2d ic oc kSize pad act).inChannels = ic
• (Layer.convBn ic oc kSize stride pad).inChannels = ic
• (Layer.layerNorm d).inChannels = d
• (Layer.dense fi fanOut act).inChannels = fi
• (Layer.residualBlock ic oc nBlocks firstStride).inChannels = ic
• (Layer.bottleneckBlock ic oc nBlocks firstStride).inChannels = ic
• (Layer.separableConv ic oc stride).inChannels = ic
• (Layer.invertedResidual ic oc expand stride nBlocks).inChannels = ic
• (Layer.mbConv ic oc expand kSize stride nBlocks useSE act).inChannels = ic
• (Layer.mbConvV3 ic oc expandCh kSize stride useSE act).inChannels = ic
• (Layer.fusedMbConv ic oc expand kSize stride nBlocks useSE).inChannels = ic
• (Layer.uib ic oc expand stride preDWk postDWk).inChannels = ic
• (Layer.fireModule ic squeeze expand1x1 expand3x3).inChannels = ic
• (Layer.patchEmbed ic dim patchSize nPatches).inChannels = ic
• (Layer.transformerEncoder dim heads mlpDim nBlocks causalMask keepSequence flashAttn rope).inChannels = dim
• (Layer.mambaBlock dim stateSize expand nBlocks).inChannels = dim
• (Layer.swinStage dim heads mlpDim windowSize nBlocks).inChannels = dim
• (Layer.patchMerging inDim outDim).inChannels = inDim
• (Layer.unetDown ic oc).inChannels = ic
• (Layer.unetUp ic oc).inChannels = ic
• (Layer.transformerDecoder dim heads mlpDim nBlocks nQueries).inChannels = dim
• (Layer.detrHeads dim nClasses).inChannels = dim
• (Layer.shuffleBlock ic oc groups nUnits).inChannels = ic
• (Layer.shuffleV2Block ic oc nUnits).inChannels = ic
• (Layer.asppModule ic oc).inChannels = ic
• (Layer.fpnModule c2 c3 c4 c5 target).inChannels = c5
• (Layer.fpnDetect oc c3 c4 c5 g5 A tower).inChannels = c5
• (Layer.evoformerBlock msaCh pairChannels nBlocks).inChannels = msaCh
• (Layer.structureModule sCh pairChannels nBlocks).inChannels = sCh
• (Layer.mobileVitBlock ic dim heads mlpDim nTxBlocks).inChannels = ic
• (Layer.convNextStage c nBlocks norm act).inChannels = c
• (Layer.convNextDownsample ic oc norm).inChannels = ic
• (Layer.convNextStem ic oc patch).inChannels = ic
• (Layer.waveNetBlock residualCh skipCh nLayers).inChannels = residualCh
• (Layer.positionalEncoding inputDim numFrequencies).inChannels = inputDim
• (Layer.nerfMLP encodedPosDim encodedDirDim hiddenDim).inChannels = encodedPosDim
• (Layer.darknetBlock c nBlocks).inChannels = c
• (Layer.cspBlock ic oc nBlocks).inChannels = ic
• (Layer.inceptionModule ic b1out b2reduce b2out b3reduce b3out b4out).inChannels = ic
• (Layer.denseBlock ic growthRate nLayers).inChannels = ic
• (Layer.transitionLayer ic oc).inChannels = ic
• (Layer.tokenPositionEmbed v t dModel ids gather posEmb).inChannels = if (ids || gather) = true then t else v * t
• (Layer.lmHead d vocabSize seqLen).inChannels = d
• Layer.spatialFlatten.inChannels = 0
• (Layer.spatialUnflatten c height width).inChannels = c
• x✝.inChannels = 0

def NetSpec.validate (s : NetSpec) : Option String

Validate that channel dimensions chain correctly through the spec. Returns none if valid, or some errorMessage describing the first mismatch.

Equations

• One or more equations did not get rendered due to their size.

def NetSpec.validate! (s : NetSpec) : IO Unit

Validate at build time. Use as: #eval mySpec.validate!

Equations

• s.validate! = match s.validate with | none => pure () | some err => throw (IO.userError (toString s.name ++ toString ": " ++ toString err))

inductive SummarySize : Type

Which quantity the summary's first line reports. Bestiary entries cover images, token streams, residue chains and raw audio, and a few (NeRF) have no meaningful input extent.

• image : SummarySize
• tokens : SummarySize
• residues : SummarySize
• samples : SummarySize
• omitted : SummarySize

@[implicit_reducible]

instance instReprSummarySize : Repr SummarySize

Equations

• instReprSummarySize = { reprPrec := instReprSummarySize.repr }

def instReprSummarySize.repr : SummarySize → Nat → Std.Format

Equations

• instReprSummarySize.repr SummarySize.image prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SummarySize.image")).group prec✝
• instReprSummarySize.repr SummarySize.tokens prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SummarySize.tokens")).group prec✝
• instReprSummarySize.repr SummarySize.residues prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SummarySize.residues")).group prec✝
• instReprSummarySize.repr SummarySize.samples prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SummarySize.samples")).group prec✝
• instReprSummarySize.repr SummarySize.omitted prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SummarySize.omitted")).group prec✝

@[implicit_reducible]

instance instDecidableEqSummarySize : DecidableEq SummarySize

Equations

• instDecidableEqSummarySize x✝ y✝ = if h : x✝.ctorIdx = y✝.ctorIdx then isTrue ⋯ else isFalse ⋯

inductive ParamUnit : Type

How params is abbreviated after the exact count. bare prints the count alone.

• millions : ParamUnit
• thousands : ParamUnit
• bare : ParamUnit

@[implicit_reducible]

instance instReprParamUnit : Repr ParamUnit

Equations

• instReprParamUnit = { reprPrec := instReprParamUnit.repr }

def instReprParamUnit.repr : ParamUnit → Nat → Std.Format

Equations

• instReprParamUnit.repr ParamUnit.millions prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "ParamUnit.millions")).group prec✝
• instReprParamUnit.repr ParamUnit.thousands prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "ParamUnit.thousands")).group prec✝
• instReprParamUnit.repr ParamUnit.bare prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "ParamUnit.bare")).group prec✝

@[implicit_reducible]

instance instDecidableEqParamUnit : DecidableEq ParamUnit

Equations

• instDecidableEqParamUnit x✝ y✝ = if h : x✝.ctorIdx = y✝.ctorIdx then isTrue ⋯ else isFalse ⋯

def NetSpec.summarize (s : NetSpec) (size : SummarySize := SummarySize.image) (unit : ParamUnit := ParamUnit.millions) (okNote : String := "") : IO Unit

The shared Bestiary summary block. Every Bestiary/*.lean entry printed its own private def summarize; all 41 were the same ten lines varying only in the first line's label, the params abbreviation and an optional parenthetical on the validate : OK line.

⚠ The ── {name} ── and params<pad>: {n} lines are load-bearing, not decoration: tests/test_bestiary_params.py parses exactly those two shapes out of each binary's stdout to pin all 189 variants against tests/bestiary_params.yml (the CI drift guard in .github/workflows/certs.yml). Keep both shapes if you touch this.

Equations

• One or more equations did not get rendered due to their size.