ML specification types: Layer, NetSpec, TrainConfig, DatasetKind.

inductive Activation : Type

• relu : Activation
• relu6 : Activation
• identity : Activation
• swish : Activation
• hSwish : Activation
• gelu : Activation

@[implicit_reducible]

instance instReprActivation : Repr Activation

Equations

• instReprActivation = { reprPrec := instReprActivation.repr }

def instReprActivation.repr : Activation → Nat → Std.Format

Equations

• instReprActivation.repr Activation.relu prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Activation.relu")).group prec✝
• instReprActivation.repr Activation.relu6 prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Activation.relu6")).group prec✝
• instReprActivation.repr Activation.identity prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Activation.identity")).group prec✝
• instReprActivation.repr Activation.swish prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Activation.swish")).group prec✝
• instReprActivation.repr Activation.hSwish prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Activation.hSwish")).group prec✝
• instReprActivation.repr Activation.gelu prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Activation.gelu")).group prec✝

@[implicit_reducible]

instance instBEqActivation : BEq Activation

Equations

• instBEqActivation = { beq := instBEqActivation.beq }

def instBEqActivation.beq : Activation → Activation → Bool

Equations

• instBEqActivation.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

inductive Normalization : Type

Normalization choice — picked at the block level by primitives that can run with either LayerNorm or BatchNorm (e.g. ConvNeXt, in its original LN form or a hypothetical BN variant for ablation).

• bn : Normalization
• ln : Normalization

@[implicit_reducible]

instance instReprNormalization : Repr Normalization

Equations

• instReprNormalization = { reprPrec := instReprNormalization.repr }

def instReprNormalization.repr : Normalization → Nat → Std.Format

Equations

• instReprNormalization.repr Normalization.bn prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Normalization.bn")).group prec✝
• instReprNormalization.repr Normalization.ln prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Normalization.ln")).group prec✝

@[implicit_reducible]

instance instBEqNormalization : BEq Normalization

Equations

• instBEqNormalization = { beq := instBEqNormalization.beq }

def instBEqNormalization.beq : Normalization → Normalization → Bool

Equations

• instBEqNormalization.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

inductive Padding : Type

• same : Padding
• valid : Padding

def instReprPadding.repr : Padding → Nat → Std.Format

Equations

• instReprPadding.repr Padding.same prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Padding.same")).group prec✝
• instReprPadding.repr Padding.valid prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Padding.valid")).group prec✝

@[implicit_reducible]

instance instReprPadding : Repr Padding

Equations

• instReprPadding = { reprPrec := instReprPadding.repr }

@[implicit_reducible]

instance instBEqPadding : BEq Padding

Equations

• instBEqPadding = { beq := instBEqPadding.beq }

def instBEqPadding.beq : Padding → Padding → Bool

Equations

• instBEqPadding.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

inductive PadStyle : Type

How a layer's .same padding is spelled when it is STRIDED.

.same states an intent ("keep the output size"), and at stride 1 with an odd kernel the two spellings below are bit-identical — which is every conv in the kit except the strided ones. They differ only when strided on an even input, where the two grids sit one input position apart:

• .xlaSame — XLA 'SAME', extra row on the HIGH side. 7×7/s2 on 224 pads (2,3); 3×3/s2 on an even input pads (0,1). This is the reference for the TF-origin ports (MobileNetV2/V4, EfficientNet), where asymmetric 'SAME' genuinely IS the published net.
• .symmetric — torchvision / He et al., (k-1)//2 on both sides, i.e. nn.Conv2d(padding=k//2). This is the reference for the ResNet family.

⚠ This exists for the same reason NetSpec.convBnAct does: the JAX emitter stated the convention independently of the render and had no way to say anything but 'SAME', so the ResNet references silently used XLA padding at their 7×7/s2 stem while the verified render used torchvision's symmetric pad. Neither shape nor op count nor arity can see the difference (planning/archive/mnv4_verified.md §3c/§4b).

⚠ The Python helpers conv2d/conv_bn already DEFAULT to symmetric — that was fixed 2026-08-04 — but the top-level layer emitter passed an explicit padding='SAME' that overrode the default at exactly the stem. .symmetric here means "pass nothing and let the helper's default apply", so there is one statement of the convention, not two.

• xlaSame : PadStyle
• symmetric : PadStyle

@[implicit_reducible]

instance instReprPadStyle : Repr PadStyle

Equations

• instReprPadStyle = { reprPrec := instReprPadStyle.repr }

def instReprPadStyle.repr : PadStyle → Nat → Std.Format

Equations

• instReprPadStyle.repr PadStyle.xlaSame prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "PadStyle.xlaSame")).group prec✝
• instReprPadStyle.repr PadStyle.symmetric prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "PadStyle.symmetric")).group prec✝

def instBEqPadStyle.beq : PadStyle → PadStyle → Bool

Equations

• instBEqPadStyle.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

@[implicit_reducible]

instance instBEqPadStyle : BEq PadStyle

Equations

• instBEqPadStyle = { beq := instBEqPadStyle.beq }

inductive Layer : Type

• conv2d (ic oc kSize : Nat) (pad : Padding) (act : Activation) : Layer
• convBn (ic oc kSize stride : Nat) (pad : Padding) : Layer
• maxPool (size stride : Nat) : Layer
• globalAvgPool : Layer
• flatten : Layer
• layerNorm (dim : Nat) : Layer
• dense (fanIn fanOut : Nat) (act : Activation) : Layer
• residualBlock (ic oc nBlocks firstStride : Nat) : Layer
• bottleneckBlock (ic oc nBlocks firstStride : Nat) : Layer
• separableConv (ic oc stride : Nat) : Layer
• invertedResidual (ic oc expand stride nBlocks : Nat) : Layer
• mbConv (ic oc expand kSize stride nBlocks : Nat) (useSE : Bool) (act : Activation := Activation.swish) : Layer
• mbConvV3 (ic oc expandCh kSize stride : Nat) (useSE : Bool) (act : Activation := Activation.relu) : Layer
• fusedMbConv (ic oc expand kSize stride nBlocks : Nat) (useSE : Bool) : Layer
• uib (ic oc expand stride preDWk postDWk : Nat) : Layer
• fireModule (ic squeeze expand1x1 expand3x3 : Nat) : Layer
• patchEmbed (ic dim patchSize nPatches : Nat) : Layer
• transformerEncoder (dim heads mlpDim nBlocks : Nat) (causalMask keepSequence flashAttn rope : Bool := false) : Layer
• mambaBlock (dim stateSize expand nBlocks : Nat) : Layer
• swinStage (dim heads mlpDim windowSize nBlocks : Nat) : Layer
• patchMerging (inDim outDim : Nat) : Layer
• unetDown (ic oc : Nat) : Layer
• unetUp (ic oc : Nat) : Layer
• bilinearUpsample (scale : Nat) : Layer
• transformerDecoder (dim heads mlpDim nBlocks nQueries : Nat) : Layer
• detrHeads (dim nClasses : Nat) : Layer
• shuffleBlock (ic oc groups nUnits : Nat) : Layer
• shuffleV2Block (ic oc nUnits : Nat) : Layer
• evoformerBlock (msaChannels pairChannels nBlocks : Nat) : Layer
• structureModule (singleChannels pairChannels nBlocks : Nat) : Layer
• mobileVitBlock (ic dim heads mlpDim nTxBlocks : Nat) : Layer
• convNextStage (channels nBlocks : Nat) (norm : Normalization := Normalization.ln) (act : Activation := Activation.gelu) : Layer
• convNextDownsample (ic oc : Nat) (norm : Normalization := Normalization.ln) : Layer
• convNextStem (ic oc patch : Nat) : Layer
• waveNetBlock (residualCh skipCh nLayers : Nat) : Layer
• positionalEncoding (inputDim numFrequencies : Nat) : Layer
• nerfMLP (encodedPosDim encodedDirDim hiddenDim : Nat) : Layer
• darknetBlock (channels nBlocks : Nat) : Layer
• cspBlock (ic oc nBlocks : Nat) : Layer
• inceptionModule (ic b1out b2reduce b2out b3reduce b3out b4out : Nat) : Layer
• asppModule (ic oc : Nat) : Layer
• fpnModule (c2 c3 c4 c5 target : Nat) : Layer
• fpnDetect (oc c3 c4 c5 g5 A tower : Nat) : Layer
• denseBlock (ic growthRate nLayers : Nat) : Layer
• transitionLayer (ic oc : Nat) : Layer
• spatialFlatten : Layer
• spatialUnflatten (channels height width : Nat) : Layer
• tokenPositionEmbed (vocabSize seqLen dModel : Nat) (idsInput gather : Bool := false) (posEmb : Bool := true) : Layer
• lmHead (dModel vocabSize seqLen : Nat) : Layer
• timeCondAdd (channels nFreq : Nat) : Layer

def instReprLayer.repr : Layer → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.
• instReprLayer.repr Layer.globalAvgPool prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Layer.globalAvgPool")).group prec✝
• instReprLayer.repr Layer.flatten prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Layer.flatten")).group prec✝
• instReprLayer.repr Layer.spatialFlatten prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Layer.spatialFlatten")).group prec✝

@[implicit_reducible]

instance instReprLayer : Repr Layer

Equations

• instReprLayer = { reprPrec := instReprLayer.repr }

structure NetSpec : Type

• name : String
• layers : List Layer
• imageH : Nat
• imageW : Nat
• detStride : Nat

  Total downsampling stride from input to the detection feature map, for YOLO-style specs: the grid is imageH / detStride. 32 for the standard ResNet stride-32 backbone (grid = 224/32 = 7, 448/32 = 14); set to 16 for a stride-16 tap (last block stride 1) to double the grid (448/16 = 28). Only consulted on the detection path; irrelevant to classifiers/seg.

• buildTag : String

  Optional suffix on buildPrefix, so two runs of the same architecture under different training configs get their own params, MLIR, and vmfb instead of overwriting each other.

  buildPrefix is derived from name alone, which is right — the name describes the architecture, and the loss is not part of the architecture. But it means a loss ablation is a set of runs that all claim the same .lake/build/<net>_params.bin and the same <net>_train_step.vmfb. Run them sequentially and each silently clobbers the last (which is why brats-predict grew an explicit params override); run them concurrently, one per GPU, and they race on the vmfb mid-compile.

  Set this to the arm's name and the ablation parallelizes cleanly. Empty (the default) reproduces the old paths exactly, so no existing trainer moves.

• convBnAct : Activation

  Activation applied after a .convBn layer, per net.

  ⚠ This exists because the JAX emitter hardcoded jax.nn.relu(x) after every .convBn and had no way to say otherwise — which silently made two references deviate from the nets they are supposed to be:

  • MobileNetV2 is ReLU6 throughout (stem and head included), not ReLU.
  • EfficientNet-B0 is SiLU/swish throughout, not ReLU.

  In both cases the verified render was already paper-faithful and the reference was wrong, so this moves the reference onto the render rather than the other way round (planning/archive/mnv4_verified.md §3f/§3h).

  ⚠ The MNv2 deviation is inert on small activations — relu6 ≡ relu below 6 — which is exactly why every forward tie at --scale 0.1 passed with it present. It only appears under He-scaled weights. Do not treat "the tie passed" as evidence the activations agree.

  Scoped PER NET, not per layer, because every net in the kit uses one activation at all of its .convBn sites. A net that genuinely mixed them would need a Layer field instead; none does, and 458 construction sites is the price of finding out.

• convPadStyle : PadStyle

  How .same is spelled at a strided top-level .conv2d/.convBn, per net.

  See PadStyle. Defaults to .xlaSame, which is what the emitter has always emitted, so every net that does not set this stays byte-identical. The ResNet-family mains set .symmetric because torchvision's Conv2d(3,64,7,stride=2,padding=3) is the net their verified render already implements.

  ⚠ Only bites at a STRIDED top-level conv. Block-internal convs (basic_block_down, bottleneck_block_down, …) call the helpers without a padding argument and so already take the symmetric default; the mobile/TF-origin blocks pass 'SAME' explicitly inside their own helpers and are likewise untouched by this field. It moves exactly the stem, which is the one site the 2026-08-04 helper-default fix could not reach.

  Scoped per net rather than per layer for the same reason as convBnAct: no net in the kit mixes the two conventions across its own top-level convs.

def instReprNetSpec.repr : NetSpec → Nat → Std.Format

Equations

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

@[implicit_reducible]

instance instReprNetSpec : Repr NetSpec

Equations

• instReprNetSpec = { reprPrec := instReprNetSpec.repr }

inductive LossKind : Type

The "kind" of supervised loss the train step computes. Picks the label-tensor shape + the forward/backward formula. Modifiers like useFocal, labelSmoothing, and the aug flags (useMixup etc.) layer on top — LossKind only captures the primary loss shape.

Used by compileVmfbs for a single-match mutex check and to drive the codegen flag set. Defaults to .classCE for back-compat with every existing trainer; if left at the default, compileVmfbs derives the effective kind from the older booleans (useYolov1, useSeg, useMixup/useCutmix/useKnnMixup) so callers don't have to update.

See planning/archive/yolo_final.md Refactor R1 for the motivation.

• classCE : LossKind

  Default: int32 [B] class label, softmax cross-entropy. Compatible with useFocal (focal modifier) and labelSmoothing.

• softLabelCE : LossKind

  Float [B, NC] soft labels (mixup/cutmix/knn-mixup output). Compatible with labelSmoothing (already baked in by the caller).

• perPixelCE : LossKind

  Int32 [B, H, W] per-pixel label tensor (segmentation). Phase 0 of the UNet demo — see planning/archive/unet_demo.md.

• perPixelDice : LossKind

  Soft Dice over the softmax probabilities, int32 [B, H, W] labels — same ABI as perPixelCE, different loss block. Dice is computed per-class over the whole batch and meaned: 1 - mean_c (2·Σ p_c·y_c + ε) / (Σ p_c + Σ y_c + ε).

  The point of it: per-pixel CE is a mean over pixels, so a class occupying 0.5% of pixels contributes 0.5% of the loss and the cheapest descent direction is to predict it away. Dice is a ratio per class, so every class carries equal weight no matter how few pixels it owns. See planning/archive/brats_demo.md — on BraTS, CE collapses all three tumour classes to IoU 0 (mIoU 0.243 ≈ the trivial background-only predictor).

  Batch-Dice (reducing over B as well as H,W) rather than per-sample: with rare classes and small batches a sample may contain none of a class at all, which makes per-sample Dice for it degenerate.

• perPixelDiceCE : LossKind

  perPixelDice + perPixelCE, summed (loss and gradient both). The standard medical-segmentation default: Dice supplies the class-balanced signal, CE regularizes it (pure Dice has a noisy gradient early, when the softmax is near-uniform and every denominator is large).

• perPixelWeightedCE (weights : List Float) : LossKind

  Class-weighted per-pixel softmax CE: weights[c] scales the loss and the gradient of every pixel whose true class is c. Same ABI as perPixelCE; weights.length must equal the class count.

  This is the lever perPixelDice was supposed to be and isn't. Dice's gradient carries a factor of p_i from the softmax Jacobian, so it vanishes exactly where a collapsed class needs rescuing — measured on BraTS at 0.02% of CE's gradient once p₃ ≈ 2e-5 (scripts/seg_dice_vanishing_grad_probe.py). CE's seed is (p - y)/N, which is -1/N at p = 0: flat, and wholly indifferent to the collapse. Scaling that by w_c therefore keeps a live signal all the way down, which is precisely what Dice cannot do. See planning/archive/brats_demo.md Workstream B'.

  Reduction is the weighted mean Σ_k w_{y_k}·CE_k / Σ_k w_{y_k} (torch's CrossEntropyLoss(weight=…, reduction='mean') semantics), not /N. Two reasons, both practical: the loss stays on the same scale as unweighted CE so the arms of an ablation are readable against each other, and — since both sums are linear in w — the loss is invariant to the overall scale of weights. Only the ratios matter, so a caller cannot accidentally change the effective learning rate by normalizing its weight vector differently. Σ_k w_{y_k} depends only on the labels, so it is a constant w.r.t. the logits and contributes no gradient term.

• perPixelFocalCE (gamma : Float) : LossKind

  Per-pixel focal CE (Lin et al., RetinaNet): -(1-p_t)^γ · log p_t, meaned over pixels. Same ABI as perPixelCE. γ = 0 is exactly perPixelCE.

  The third distinct answer to the imbalance, and mechanically the opposite of perPixelWeightedCE — worth stating, because "focal and class weights both reweight the loss" hides the whole point:

  • weighted CE amplifies the rare class. A static, per-class factor from the label frequencies.
  • focal suppresses the easy class. A dynamic, per-pixel factor from the current prediction. At p_t → 1 (confident background — 97% of BraTS) the (1-p_t)^γ factor crushes the gradient toward 0. At p_t → 0 it tends to 1 and the gradient tends to CE's: focal does not amplify the collapsed class, it defunds the majority drowning it out.

  So focal needs no frequency statistics and cannot be mis-tuned by a bad weight vector, but it also cannot help a class that is rare and easy.

  α is deliberately omitted. The paper's α_t is a per-class weight, i.e. exactly perPixelWeightedCE's mechanism — folding it in here would confound the two arms of the very ablation this exists for. Compose them later, once each is understood alone.

  NB the gradient here is the true one, including the derivative of the (1-p_t)^γ factor. Contrast the YOLOv1 objectness path, which detaches that weight (MlirCodegen.lean, %y1f_w0) — a defensible approximation, but one whose d_logits is not the derivative of any loss it states, and so could not be FD-verified the way this is.

• floatTargetMse : LossKind

  Float [B, C, H, W] target tensor with per-pixel MSE (DDPM, autoencoder regression). Caller passes ddpmOutShape.

• yolov1Masked : LossKind

  YOLOv1: float [B, perCell, gridH, gridW] target + float [B, gridH, gridW] per-cell mask. 5-term masked MSE with √ ε-floor on the box-dim terms (see planning/archive/yolo_final.md Phase 1).

• bce : LossKind

  Binary cross-entropy with logits over multi-hot [B, NC] targets — timm "ResNet Strikes Back" RSB-A2's loss. Each class is an independent sigmoid; the mixup/cutmix soft-label path produces the [B,NC] target directly (hard labels are one-hot'd, with optional label smoothing). JAX-only (the IREE/MLIR backend does not implement it). Reduction is timm's mean over B×C. See planning/archive/rsb_a2_resnet50.md.

@[implicit_reducible]

instance instReprLossKind : Repr LossKind

Equations

• instReprLossKind = { reprPrec := instReprLossKind.repr }

def instReprLossKind.repr : LossKind → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.
• instReprLossKind.repr LossKind.classCE prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.classCE")).group prec✝
• instReprLossKind.repr LossKind.softLabelCE prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.softLabelCE")).group prec✝
• instReprLossKind.repr LossKind.perPixelCE prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.perPixelCE")).group prec✝
• instReprLossKind.repr LossKind.perPixelDice prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.perPixelDice")).group prec✝
• instReprLossKind.repr LossKind.perPixelDiceCE prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.perPixelDiceCE")).group prec✝
• instReprLossKind.repr LossKind.floatTargetMse prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.floatTargetMse")).group prec✝
• instReprLossKind.repr LossKind.yolov1Masked prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.yolov1Masked")).group prec✝
• instReprLossKind.repr LossKind.bce prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "LossKind.bce")).group prec✝

@[implicit_reducible]

instance instBEqLossKind : BEq LossKind

Equations

• instBEqLossKind = { beq := instBEqLossKind.beq }

def instBEqLossKind.beq : LossKind → LossKind → Bool

Equations

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

inductive SegLoss : Type

Which loss block the segmentation path emits. All three share one ABI (int32 [B,H,W] labels, the trainStepAdamF32Seg dispatch, the mIoU eval harness) — only the emitted loss + gradient differ, so this rides alongside useSeg rather than replacing it.

• ce : SegLoss

  Per-pixel softmax cross-entropy (the original UNet-demo path).

• dice : SegLoss

  Soft Dice only.

• diceCE : SegLoss

  Dice + CE, summed.

• weightedCE (weights : List Float) : SegLoss

  Per-pixel CE with a per-class weight on the true class. See LossKind.perPixelWeightedCE for the semantics and the argument.

• focalCE (gamma : Float) : SegLoss

  Per-pixel focal CE, -(1-p_t)^γ·log p_t. See LossKind.perPixelFocalCE.

def instReprSegLoss.repr : SegLoss → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.
• instReprSegLoss.repr SegLoss.ce prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SegLoss.ce")).group prec✝
• instReprSegLoss.repr SegLoss.dice prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SegLoss.dice")).group prec✝
• instReprSegLoss.repr SegLoss.diceCE prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "SegLoss.diceCE")).group prec✝

@[implicit_reducible]

instance instReprSegLoss : Repr SegLoss

Equations

• instReprSegLoss = { reprPrec := instReprSegLoss.repr }

@[implicit_reducible]

instance instBEqSegLoss : BEq SegLoss

Equations

• instBEqSegLoss = { beq := instBEqSegLoss.beq }

def instBEqSegLoss.beq : SegLoss → SegLoss → Bool

Equations

• instBEqSegLoss.beq SegLoss.ce SegLoss.ce = true
• instBEqSegLoss.beq SegLoss.dice SegLoss.dice = true
• instBEqSegLoss.beq SegLoss.diceCE SegLoss.diceCE = true
• instBEqSegLoss.beq (SegLoss.weightedCE a) (SegLoss.weightedCE b) = (a == b)
• instBEqSegLoss.beq (SegLoss.focalCE a) (SegLoss.focalCE b) = (a == b)
• instBEqSegLoss.beq x✝¹ x✝ = false

@[implicit_reducible]

instance instInhabitedSegLoss : Inhabited SegLoss

Equations

• instInhabitedSegLoss = { default := instInhabitedSegLoss.default }

def LossKind.isSeg : LossKind → Bool

Does this loss run on the segmentation path? Every per-pixel kind shares the int32 [B,H,W] label ABI and the seg train-step dispatch.

Single source of truth: compileVmfbs, NetSpec.train, and runTraining each need this answer, and deriving it three times independently is how they drift apart.

Equations

• LossKind.perPixelCE.isSeg = true
• LossKind.perPixelDice.isSeg = true
• LossKind.perPixelDiceCE.isSeg = true
• (LossKind.perPixelWeightedCE weights).isSeg = true
• (LossKind.perPixelFocalCE gamma).isSeg = true
• x✝.isSeg = false

def LossKind.segLoss : LossKind → SegLoss

Which seg loss block to emit. Non-seg kinds answer .ce and are never asked (the codegen only consults this under useSeg).

Equations

• LossKind.perPixelDice.segLoss = SegLoss.dice
• LossKind.perPixelDiceCE.segLoss = SegLoss.diceCE
• (LossKind.perPixelWeightedCE w).segLoss = SegLoss.weightedCE w
• (LossKind.perPixelFocalCE g).segLoss = SegLoss.focalCE g
• x✝.segLoss = SegLoss.ce

inductive OptimizerKind : Type

Optimizer selector for the training loop. Added additively over the legacy TrainConfig.useAdam bool (à la LossKind over the older loss booleans): the JAX backend derives the effective optimizer from useAdam when this is left at the .sgd default, so no existing config needs to change.

• sgd : OptimizerKind

  Plain SGD, or SGD + heavy-ball momentum when TrainConfig.momentum > 0.

• adam : OptimizerKind

  Adam / AdamW (decoupled weight decay when weightDecay > 0).

• rmsprop : OptimizerKind

  RMSprop with momentum — the native MobileNetV2 / EfficientNet optimizer. v = ρ·v + (1-ρ)·g²; buf = μ·buf + g/(√v+ε); p -= lr·buf, with ρ = rmspropDecay, μ = momentum, ε = rmspropEps. Weight decay stays coupled into the gradient (the form those papers use), unlike AdamW.

• lamb : OptimizerKind

  LAMB (You et al. 2019) — the large-batch optimizer in timm's "ResNet Strikes Back" RSB-A2 recipe. Adam moments (m, v, t) form the per-param direction r = m̂/(√v̂+ε) + λ·θ (DECOUPLED weight decay λ = weightDecay folded into the direction), then a layer-wise trust ratio ‖θ‖ / ‖r‖ rescales the step: θ -= lr · (‖θ‖/‖r‖) · r. The trust ratio is 1.0 wherever ‖θ‖ or ‖r‖ is 0 (timm convention). β1=0.9, β2=0.999, ε=1e-6. opt_state shape matches .adam: (m, v, t).

• muon : OptimizerKind

  Muon (Jordan 2024) — MomentUm Orthogonalized by Newton–Schulz. The heavy-ball momentum buffer is polar-projected onto the (semi-)orthogonal matrices (G = UΣVᵀ ↦ UVᵀ) by a fixed 5-step Newton–Schulz iteration (pure matmul, no SVD), so every singular direction gets an equal-size step. Applies ONLY to 2D weight matrices; non-2D params (biases, norms, embeddings, small heads) fall back to AdamW. IREE/MLIR perf path reads TrainConfig.useMuon. UNVERIFIED. See planning/archive/muon.md.

• shampoo : OptimizerKind

  Shampoo (Gupta–Koren–Singer 2018) — Kronecker-factored full-matrix preconditioner. For a 2D weight W∈ℝ^{m×n} it accumulates L=Σ GGᵀ (m×m) and R=Σ GᵀG (n×n) and steps W -= η·L^{-1/4}·G·R^{-1/4}. The inverse 4th roots are computed matmul-only by trace-scaled coupled-Newton inverse-sqrt (reuses Muon's NS machinery). Single-step (un-accumulated) Shampoo IS Muon = the polar factor UVᵀ — the demo's jewel. Demo scope: applies ONLY to square 2D weight matrices (m==n), where L[m,m] and R[n,n] fit exactly into the existing m/v optimizer slots — so like Muon the train-step signature is Adam-identical (no extra state buffers) and it drives through the existing Adam FFI. Non-square 2D weights and non-2D params fall back to AdamW. L/R use EMA accumulation and are εI- regularized at inversion time (state slots init to 0, so no host change). IREE/MLIR perf path reads TrainConfig.useShampoo. UNVERIFIED. See planning/archive/shampoo.md.

def instReprOptimizerKind.repr : OptimizerKind → Nat → Std.Format

Equations

• instReprOptimizerKind.repr OptimizerKind.sgd prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "OptimizerKind.sgd")).group prec✝
• instReprOptimizerKind.repr OptimizerKind.adam prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "OptimizerKind.adam")).group prec✝
• instReprOptimizerKind.repr OptimizerKind.rmsprop prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "OptimizerKind.rmsprop")).group prec✝
• instReprOptimizerKind.repr OptimizerKind.lamb prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "OptimizerKind.lamb")).group prec✝
• instReprOptimizerKind.repr OptimizerKind.muon prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "OptimizerKind.muon")).group prec✝
• instReprOptimizerKind.repr OptimizerKind.shampoo prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "OptimizerKind.shampoo")).group prec✝

@[implicit_reducible]

instance instReprOptimizerKind : Repr OptimizerKind

Equations

• instReprOptimizerKind = { reprPrec := instReprOptimizerKind.repr }

def instBEqOptimizerKind.beq : OptimizerKind → OptimizerKind → Bool

Equations

• instBEqOptimizerKind.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

@[implicit_reducible]

instance instBEqOptimizerKind : BEq OptimizerKind

Equations

• instBEqOptimizerKind = { beq := instBEqOptimizerKind.beq }

@[implicit_reducible]

instance instDecidableEqOptimizerKind : DecidableEq OptimizerKind

Equations

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

structure TrainConfig : Type

• learningRate : Float
• batchSize : Nat
• epochs : Nat
• seed : Nat
• momentum : Float
• useAdam : Bool
• optimizer : OptimizerKind

  Optimizer selector (additive over useAdam). Left at the .sgd default, the JAX backend derives the effective optimizer from useAdam (true → Adam) for back-compat; set explicitly to .rmsprop (or .adam) to override. The IREE/MLIR backend still reads useAdam.

• useMuon : Bool

  Muon selector for the IREE/MLIR perf path (additive over useAdam, like optimizer is for JAX). When true, every 2D weight matrix is updated by Muon (Newton–Schulz polar projection); all non-2D params use AdamW. Left false by default so no existing config changes behavior. See planning/archive/muon.md.

• useShampoo : Bool

  Shampoo selector for the IREE/MLIR perf path (additive over useAdam, like useMuon). When true, every square 2D weight matrix (m==n, both dims ≥ 16) is updated by Shampoo (Kronecker L^{-1/4}·G·R^{-1/4}); non-square 2D weights and all non-2D params use AdamW. The L/R state reuses the m/v slots (square ⇒ same shape), so the module stays Adam-signature- identical. Left false by default. See planning/archive/shampoo.md.

• rmspropDecay : Float

  RMSprop running-mean-square decay ρ (only used when optimizer = .rmsprop).

• rmspropEps : Float

  RMSprop denominator ε — NOT 1e-8: MobileNetV2 uses 1.0, EfficientNet 1e-3; the large value is part of those recipes.

• weightDecay : Float
• cosineDecay : Bool
• warmupEpochs : Nat
• augment : Bool
• labelSmoothing : Float
• useFocal : Bool

  Focal loss (Lin et al. 2017): replace CE loss with -(1-p_y)^γ · log(p_y). Down-weights well-classified examples, up-weights hard ones. Restricted to the int-label path (no soft labels) and labelSmoothing must be 0 — focal mixes poorly with both. γ=2.0 is the paper default.

• focalGamma : Float
• useDiouBox : Bool

  YOLOv1 box loss: false = the published √-MSE coord terms; true = a DIoU box loss on box0 with a positive box parameterization (cx=(j+σ(tx))/gW, w=exp(tw)), the detection-infra brick #1 (planning/archive/yolo_drone.md WS-D). Only consulted on the .yolov1Masked path. NB: a DIoU-trained model must be decoded with the same σ/exp (scripts/yolo_map_visdrone.py --box-param diou).

• anchors : List (Float × Float)

  YOLO anchor priors (w_rel, h_rel) for the anchor-based detector (brick #2). Empty = single-box YOLOv1. When non-empty, the yolo loss routes to the A-anchor path (perCell = A·15, box_a = anchor_a·exp(pred)); the target/head must use the matching A·15-channel layout (preprocess_visdrone --anchors).

• fpnScales : List (Nat × List (Float × Float))

  FPN multi-scale detector (planning/archive/yolo_fpn.md bite 7). Per-scale (grid, anchors) for P3/P4/P5 (e.g. [(56, a3), (28, a4), (14, a5)]). Empty = not an FPN detector. When non-empty AND the spec ends in a .fpnDetect layer, the loss routes to emitMultiScaleYoloLoss over the [B, Ntot] head concat, and the single target input is a flat [B, Σ A·15·g_s²] block (sliced per scale in the loss).

• fpnAffineScale : Float

  Box-aware affine augmentation for the FPN path: per-image scale gain 1 + U(-1,1)·fpnAffineScale and translate U(-1,1)·fpnAffineTranslate (as a fraction of the frame), fired with probability fpnAffineProb. 0 scale gain and 0 probability (the defaults) leave the pipeline byte-identical to the HSV+hflip pack, so the existing arms stay reproducible.

  Separate from augment on purpose. augment is the committed and measured pack (HSV + hflip, worth 0.1243 → 0.1674 at 50 epochs); this is the arm that changes object SCALE, which on VisDrone is the axis the difficulty actually lives on — and therefore the one that can hurt as easily as help, since a 2–5 px object scaled down goes under P3's stride-8 resolution entirely. It must be A/B'd under its own tag, on top of augment, not instead of it.

• fpnAffineTranslate : Float
• fpnAffineProb : Float
• fpnAffineWhThrPx : Float

  Drop a transformed box whose clipped side falls under fpnAffineWhThrPx pixels, or which keeps less than fpnAffineAreaThr of its area inside the frame. Ultralytics uses 2 px and 0.1; the 2 px default is wrong here, because a large share of VisDrone GT is 2–5 px to begin with and that threshold would silently delete the classes the detector is worst at.

• fpnAffineAreaThr : Float
• yoloClsWeights : List Float

  Per-class weights for the detector's classification term (planning/archive/ yolo_fpn.md T1b). weights.length must equal the detector class count (10 for VisDrone). Empty (the default) is the unweighted path and emits byte-identical MLIR.

  Motivated by measurement, not folklore: on the unweighted e12 checkpoint the class argmax collapsed onto the two most frequent classes (car 44% + pedestrian 21% of encoded positives), leaving 5/10 classes never predicted and per-class mAP pinned at ~0.0001 — see scripts/fpn_obj_separation.py and scripts/fpn_class_freq.py. Weights depend only on the target, so they are exactly constant w.r.t. the logits and the weighted gradient stays finite-difference checkable.

  Normalize so Σ_c f_c·w_c = 1 (expected weight 1 under the GT class distribution) to keep this a pure redistribution — otherwise it silently rescales the class term against the box and objectness terms too.

• yoloClsFocalGamma : Float

  Focal γ on the CLASS term (T1c). 0 = plain (weighted) softmax-CE and emits byte-identical MLIR. 2.0 is the RetinaNet value.

  FL = −w·(1−p_t)^γ·log p_t on assigned cells, p_t the softmax probability of the true class. Distinct from focalGamma, which is the OBJECTNESS focal and applies to every cell.

  Why this lever rather than more class weighting: the FPN_CLSW ladder measured static per-class weights as net harmful (mAP none 0.1774 / sqrt 0.1771 / inv 0.1368) because a fixed constant raises rare-class recall by flooding those classes with false positives, and AP is precision-sensitive — tricycle detections went 7,419 → 31,322 against 1,045 GT. Focal down-weights EASY examples whatever their class, and the weight tracks p_t as it moves, so it cannot buy recall with a permanent precision tax.

• headPriorBias : List Float

  RetinaNet prior-bias init: initialize the head's bias to log π_c instead of zero, so the net starts predicting the class prior rather than a uniform distribution. Empty (the default) leaves the head at zero bias and changes nothing.

  Orthogonal to lossKind — it is an init, not a loss — and deliberately so: it is the natural partner of .perPixelFocalCE, whose (1-p_t)^γ factor is a no-op at a uniform softmax because there is no confidence to suppress. Prior-bias init manufactures that confidence at step 0. See NetSpec.applyHeadPriorBias for the measured size of the effect.

• detPriorPi : Float

  RetinaNet prior-bias init for the FPN detector head: initialize every objectness bias to −log((1−π)/π) so the head starts at sigmoid = π. 0.0 (the default) leaves the head biases at zero, which reproduces the biasless head exactly. Typical value 0.01.

  The sigmoid-head twin of headPriorBias above, and the Tier-2 lever the loss-breakdown measurements pointed at: objectness had signal but no dynamic range, because a bias-free 1×1 conv spends its weights manufacturing the constant background offset. See NetSpec.applyDetPriorBias.

• useMixup : Bool

  DeiT-style data augmentation knobs. Setting useMixup or useCutmix switches the train-step to the soft-label codegen path; the dataloader produces a [B, NC] smoothed soft-label tensor instead of an int32 [B] vector. mixupAlpha and cutmixAlpha control the Beta-distribution shape; the paper defaults are 0.8 and 1.0 respectively. randomErasing operates on the int-label path (no soft-label conversion needed).

• mixupAlpha : Float
• useCutmix : Bool
• cutmixAlpha : Float
• useKnnMixup : Bool

  KNN-Mixup: pair each sample with its nearest neighbor in the batch (pixel-space L2) rather than a random partner. Closer manifold mixing → harder, more realistic intermediate samples. Mutually exclusive with useMixup/useCutmix (KNN takes precedence).

• knnMixupAlpha : Float
• randomErasing : Bool
• randomErasingProb : Float
• useRandAugment : Bool

  RandAugment-Color (Cubuk et al. 2019, color subset). Applied per-image before mixup/cutmix, after crop/hflip. randAugmentN ops drawn uniformly from {identity, brightness, contrast, color, autocontrast} per image, each at magnitude randAugmentM (0–10, paper default 9). No labels touched.

• randAugmentN : Nat
• randAugmentM : Float
• randAugmentGeometric : Bool

  Upgrade useRandAugment from the color-only "lite" path to the full RandAugment(N, M) sampler over the color+GEOMETRIC op set (shear/rotate/ translate via ImageProjectiveTransformV3, shared with AutoAugment). This is what ConvNeXt's recipe wants. Only meaningful when useRandAugment is on; leaving it false keeps the back-compat color-lite path (e.g. ViT).

• randAugmentMstd : Float

  DeiT/ConvNeXt RandAugment refinements (gap D), meaningful only with the geometric sampler on. randAugmentMstd (timm mstd, DeiT uses 0.5) draws each op's magnitude from N(M, mstd) clipped to [0,10] instead of a fixed M. randAugmentInc (timm inc1) uses the increasing-severity magnitude→arg mappings: solarize/posterize flip so higher M = more distortion, and the enhancement ops center at 1.0 ± sign·scaled (random direction).

• randAugmentInc : Bool
• useAutoAugment : Bool

  AutoAugment, ImageNet learned policy (Cubuk et al. 2018) — the full 25 sub-policies, applied per-image after crop/hflip on the imagenet (tfds) path. Unlike useRandAugment (color subset only), this includes the GEOMETRIC ops (shear/rotate) via tf.raw_ops.ImageProjectiveTransformV3 (core TF — dissolves the old "tfa unavailable on tf2.21" blocker) plus the full color set (posterize/solarize/equalize/autocontrast/etc). Subsumes the color RandAugment, so leave useRandAugment off when this is on. EfficientNet's original recipe; no labels touched.

• repeatedAug : Nat

  Repeated Augmentation (Hoffer et al. 2020; timm RASampler), RSB-A2's 3×. Each image contributes repeatedAug independently-augmented copies per epoch. On the tfds path this is a stream-level flat_map(repeat K) before the augment _pp, plus a re-shuffle so the copies spread across batches — an APPROXIMATION of timm's exact index-level RASampler. steps_per_epoch is unchanged, so an epoch sees ~1/K as many unique images ×K views, per the RSB recipe. 1 disables.

• trainRes : Nat

  Train/test resolution split (RSB-A3): TRAIN at trainRes×trainRes, EVAL at the spec's imageH/imageW. 0 = no split (train and eval same resolution). The generated forward infers the square resolution from the flat input length, so the conv stack + global-avg-pool run at either size (A3 trains @160, tests @224 → ~2× cheaper per step). imagenet (tfds) path only.

• testCropRatio : Float

  Explicit test-time center-crop ratio (RSB-A3 uses 0.95). 0 = the default _IMG_SIZE/(_IMG_SIZE+32) ≈ 0.875 ratio. imagenet (tfds) path only.

• dropPath : Float

  Stochastic depth (Huang et al. 2016): drop each residual block's branch with a probability that ramps linearly from 0 to dropPath across the network's residual blocks; surviving branches are scaled by 1/keep (inverted, so inference is drop-free). 0 disables. On the JAX path this threads a per-step RNG through forward; currently wired for ConvNeXt blocks.

• wdExcludeNormBias : Bool

  AdamW no_weight_decay exclusion (timm/DeiT): when true, decoupled weight decay skips 1-D params (all biases, LayerNorm γ/β, the CLS token) and the positional embedding, decaying only ≥2-D weight matrices. Matches the ViT/DeiT reference; off keeps the legacy decay-everything behavior for the other nets. AdamW path only.

• valEveryEpochs : Nat

  Validate every N epochs (plus always the final epoch) instead of every epoch. N ≤ 1 keeps every-epoch validation (byte-identical codegen). Cuts eval wall-time on large streaming datasets where the val pass is data-loading-bound (e.g. ImageNet: ~75s/epoch rebuilding the tfds val pipeline). ImageNet-streaming main only.

• gradAccumSteps : Nat

  Gradient accumulation: run gradAccumSteps micro-batches of batchSize before each optimizer update, giving an EFFECTIVE batch of batchSize × gradAccumSteps at the peak-activation cost of ONE micro-batch. The reproducibility lever for large-batch recipes (e.g. RSB LAMB @ bs2048) on small GPUs: batchSize=512 × gradAccumSteps=4 on 4×16GB instead of an 8×A100 node. learningRate should target the EFFECTIVE batch. BatchNorm uses per-micro-batch (Ghost-BN) statistics — not identical to true large-batch BN, but a benign/beneficial variant at micro≥256. N ≤ 1 keeps the single-shot update (byte-identical codegen). ImageNet-streaming main.

• useEMA : Bool

  DeiT-style training-loop knobs that average weights for the eval checkpoint. Both can be on simultaneously; eval picks EMA when both are enabled. Storage cost: one extra nParams-sized buffer per knob; runtime cost is one F32.ema call per step (EMA) or per epoch (SWA), well below GPU step time.

• emaDecay : Float
• useSWA : Bool
• swaStartEpoch : Nat

  First epoch (zero-indexed) that contributes to the SWA average. Typical recipe: 0.75 × epochs, e.g. epoch 60 of 80.

• useSWAG : Bool

  SWAG (Maddox et al. 2019): extends SWA with diagonal Σ_diag (via running mean of θ²) plus a low-rank component built from the last swagK per-epoch deviations from the SWA mean. At eval, sample swagSamples weight vectors from N(swaMean, ½Σ_diag + ½ Σ_low), run forward each, average logits. Requires useSWA=true.

• swagK : Nat
• swagSamples : Nat
• useTTA : Bool

  TTA (test-time augmentation): at periodic eval, run ttaSamples independently-augmented forwards per batch and average the logits. Augmentations are the same dataloader pipeline used for training (e.g. random crop + hflip for Imagenette), minus the soft-label ones (mixup/cutmix) which need the label. Eval-only — no train cost, M× eval cost.

• ttaSamples : Nat
• useYolov1 : Bool

  YOLOv1 5-term masked-MSE loss. See planning/archive/yolo_final.md Phase 1 + planning/archive/yolo_final.md for integration scope. Equivalent to lossKind := .yolov1Masked; the bool form predates LossKind and is retained for back-compat.

• lossKind : LossKind

  Explicit loss-kind selector. If left at the default .classCE, compileVmfbs derives the effective kind from the older booleans (useYolov1, useSeg, soft-label augs). Set explicitly to skip the derivation path or to disambiguate borderline cases. See LossKind.

• bootstrapBackbone : Option (String × Nat)

  Bootstrap from a pretrained backbone checkpoint. When set to some (paramsPath, prefixFloats), runTraining overwrites the first prefixFloats * 4 bytes of the He-init with bytes read from paramsPath. The companion <basename>_bn_stats.bin is auto-loaded too if present (the backbone's BN running stats must match the spec's BN layer count + sizes — true for YOLOv1 loading R34 weights since both have identical backbone layers).

  Phase 4 of planning/archive/yolo_final.md. Example for YOLOv1+R34: bootstrapBackbone := some (".lake/build/resnet_34_params.bin", 21284672).

• bootstrapBackboneRange : Option (String × Nat × Nat × Nat)

  Offset-aware bootstrap, as some (paramsPath, dstOffFloats, srcOffFloats, countFloats). Same job as bootstrapBackbone but for a spec whose FIRST layer differs from the checkpoint's, so the transferable weights are no longer a prefix — see NetSpec.patchInitWithPretrainedRange.

  The motivating case is r34UnetBrats: a 4-modality MRI stem ([64,4,7,7], 12,544 floats) in front of an R34 body pretrained on 3-channel RGB ([64,3,7,7], 9,408 floats). The stem keeps its He-init — MRI is not RGB, so re-learning it is the correct behaviour, not a concession — and the remaining 21,275,264 floats of the backbone transfer intact.

  Takes precedence over bootstrapBackbone when both are set. BN running stats follow the same rule as the prefix path: loaded only on an exact size match, zeros otherwise.

• checkpointEveryNEpochs : Nat

  Save intermediate {pfx}_params_e{N}.bin and {pfx}_bn_stats_e{N}.bin snapshots every checkpointEveryNEpochs epochs. 0 disables. Default 10 = align with the eval cadence so training can be killed mid-run and resumed from the most recent checkpoint (or borrowed for downstream tasks like YOLOv1 bootstrap — see bootstrapBackbone).

• evalEveryNEpochs : Nat

  Run the validation eval every N epochs (plus always on the final epoch). 0 means final-epoch only.

  Default 10 matches the historical hardcoded cadence, which is fine for a 100-epoch classifier and actively bad for a short ablation: a 10-epoch run gets exactly ONE eval, at the very end, so a 7-hour run yields no signal until it is over. On the seg path the eval is the only instrument that can see a collapsed class — no scalar in the training log can (planning/archive/brats_demo.md Workstream A) — so flying blind is worse here than anywhere else. Set it to 1-2 for ablation arms; the eval is a forward pass over val and costs minutes against 40 min/epoch.

• bf16 : Bool

  bf16 mixed precision: cast matmul operands (dense, attention QKV / scores / output, MLP, patch embed) to bfloat16, keeping master weights, LayerNorm, softmax, and conv in fp32. Measured ~2.7-3.6× on matmul-bound nets (ViT/transformers) on gfx1100; ~no effect on conv-bound nets (bf16 conv is slower on MIOpen, so convs stay fp32). See reference_bf16_gfx1100_conv_vs_gemm.

• bf16Conv : Bool

  bf16 conv compute: additionally cast the standard conv path (conv2d / conv_bn, hence the full ResNet/VGG/CIFAR-CNN conv stack) to bfloat16, returning fp32. Independent of bf16 so the AMD/MIOpen path can keep convs in fp32 (default false) while CUDA/cuDNN — where bf16 conv is ~1.6× FASTER via tensor cores — can opt in. Only meaningful when bf16 := true. Depthwise/separable convs (MobileNet/EfficientNet) still stay fp32.

• runningBN : Bool

  Running batch-norm statistics (gap A). When true, the JAX imagenet trainer tracks per-BN-layer running mean/var (EMA of batch stats, momentum bnMomentum) threaded through forward as has_aux, and EVAL normalizes with the running stats instead of the eval batch's own — the paper-faithful behaviour. Off (default) keeps the current batch-stats-at-eval path, so every existing net is byte-identical. Currently wired for the convBn + invertedResidual path (MobileNetV2); extend the BN-threading to mbconv/basic/bottleneck blocks for the other convnets. See planning/archive/jax_imagenet_sweep.md "Gap A".

• bnMomentum : Float

  BatchNorm running-statistic decay, i.e. the weight on the OLD estimate: running = bnMomentum·running + (1−bnMomentum)·batch. Only consulted when runningBN is on (it is an eval-time statistic and touches no gradient).

  ⚠⚠ This is the TensorFlow convention, and it is the reciprocal of PyTorch's. torch.nn.BatchNorm2d(momentum=m) weights the NEW batch by m, so timm's PyTorch default momentum = 0.1 is bnMomentum = 0.9 here — a 10-step-vs-100-step averaging window, not a 10× smaller one. Getting the sense backwards silently makes the running stats 10× noisier, and it is eval-only, so no loss curve moves.

  Per-net, because the nets chase different references (audited 2026-08-30 against the pinned timm 1.0.28 in .venv-timm, create_model(...).modules()):

  net	reference	value
  R50 / R34	timm (RSB) — momentum=0.1	0.9
  EfficientNet-B0	the TF paper (77.1/93.3) — TF's decay=0.99	0.99
  MobileNetV2	the TF-slim paper (72.0) — decay=0.997 [unverified]	0.99 today
  MNv4	our own 100-ep JAX run, not a paper	0.99 today

  ⚠ timm itself runs 0.9 on every one of those nets, tf_efficientnet_b0 included: BN_MOMENTUM_TF_DEFAULT exists in _efficientnet_builder.py but get_bn_args_tf() has no caller in 1.0.28, so the tf_* ports inherit only bn_eps = 1e-3. So "0.99 is right for EfficientNet" is a statement about the original TF codebase, NOT about timm — flip a net to 0.9 only along with the reference it is being scored against.

  The default 0.99 keeps every net that does not set it byte-identical, on both the JAX emitter (Jax/Codegen.lean's _bn) and the verified host-side EMA (VerifiedTrain.lean's bnMom). Under gradient accumulation both sides compensate to bnMomentum^(1/K) per micro-batch.

• vitInit : Bool

  timm/DeiT ViT weight init, replacing the generic Xavier-uniform for transformer-shaped nets. Off by default so every existing run is byte-identical; turn it on per-recipe.

  timm's init_weights_vit_timm gives every nn.Linear (QKV, attn-out, MLP fc1/fc2, and the classifier head) trunc_normal_(std=0.02) with zero bias, and leaves the patch-embed nn.Conv2d on PyTorch's default kaiming_uniform_(a=sqrt(5)), i.e. U(±1/sqrt(fan_in)).

  NB trunc_normal_'s default bounds a=-2, b=2 are ABSOLUTE, so at std=0.02 they sit at ±100σ and the truncation never fires — it is plain normal(0, 0.02). Emitted as such, matching the CLS-token and positional-embedding init already in the emitter.

  Why it matters: Xavier scales as 1/sqrt(dim) against a fixed 0.02, so the generic path is 1.8× too wide at ViT-B, 2.6× at ViT-S and 3.6× at ViT-Ti, while the patch embed is ~6× too NARROW (it divides by the output fan dim·p·p rather than the input fan ic·p·p). See planning/archive/vit_imagenet.md item 0.

• cnxInit : Bool

  ConvNeXt paper weight init, replacing the generic Xavier-uniform for the ConvNeXt-shaped nets. Off by default; turn it on per-recipe.

  ConvNeXt's reference implementation applies trunc_normal_(std=.02) with zero bias to every nn.Conv2d AND nn.Linear (_init_weights in facebookresearch/ConvNeXt, and timm's convnext.py agrees), leaving the LayerNorms at (1, 0) and LayerScale at 1e-6 — both of which the emitter already gets right.

  ⚠ This is deliberately a SEPARATE flag from vitInit, not a shared "timmInit". The two specs disagree on the conv path: ViT leaves its patch-embed nn.Conv2d on PyTorch's default U(±1/sqrt(fan_in)), while ConvNeXt trunc-normals its convs like everything else. One boolean cannot express both, and a flag named for the vendor rather than the distribution would invite exactly the misapplication below.

  ⛔ Do NOT extend either flag to the ResNet/MobileNet/EfficientNet family. "timm init" is not one thing: those nets use kaiming_normal_(mode='fan_out', nonlinearity='relu'), and the emitter's emitConvBnInit — uniform ±sqrt(6/(oc·k²)), i.e. std sqrt(2/fan_out) — is ALREADY on that scale, differing only in distribution shape. Giving a ResNet conv trunc_normal(0.02) would be a regression, not a fix.

  Why it matters: the generic emitConvBiasInit is Xavier over ic·kh·kw + oc, so ConvNeXt-T's stem lands at std 0.118 against the paper's 0.02 — 5.9× too wide, the same failure vitInit fixes for ViT and slightly worse. See planning/archive/vit_imagenet.md item 0 for the ViT half of the story.

• expLRDecayRate : Float

  Exponential LR decay (gap B), the EfficientNet/MobileNet schedule: after warmup, lr = LR · rate^((epoch − warmup) / decayEpochs). 0 = off (use cosine). EfficientNet: rate 0.97, decayEpochs 2.4; MobileNetV2: rate 0.98, decayEpochs 1.0. Selected over cosine when > 0.

• expLRDecayEpochs : Float
• dropout : Float

  Classifier dropout (gap C): dropout rate applied before the final dense head during training (inverted, scaled by 1/keep so eval is drop-free). 0 = off. EfficientNet-B0 / MobileNetV2 use 0.2. Threaded via the same drop_key as stochastic depth; requires the running-BN training flag (or drop_key≠None) to gate train-vs-eval.

• gradClipNorm : Float

  Clip gradients by global L2 norm before the optimizer step. 0 = off. DeiT default 1.0 — essential for stable ViT-from-scratch training: it lets you use the proper ~1e-3 LR without the collapse-to-chance seen at higher LR with no clipping. See planning/archive/vit_imagenet.md.

• headLrMult : Float

  Per-group LR multiplier for the (from-scratch) dense head, relative to the base LR used by the pretrained conv backbone. 1.0 = uniform LR. Used for bootstrap fine-tuning where the He-init head must learn input- dependence far faster than the backbone should drift — a single LR can't do both (head under-trains → collapse-to-marginal; raise it globally and the backbone destabilizes). e.g. YOLOv1 detection uses ~10.

@[implicit_reducible]

instance instReprTrainConfig : Repr TrainConfig

Equations

• instReprTrainConfig = { reprPrec := instReprTrainConfig.repr }

def instReprTrainConfig.repr : TrainConfig → Nat → Std.Format

Equations

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

inductive DatasetKind : Type

• mnist : DatasetKind
• cifar10 : DatasetKind
• imagenette : DatasetKind
• pets : DatasetKind
• imagenet : DatasetKind
• petsDet : DatasetKind

  YOLOv1 detection on Oxford-IIIT Pets (cat/dog head boxes, tiled into 2×2 mosaics). Images are 224×224×3 (resized at preprocess time, ImageNet-normalized on Lean read). Labels carry the YOLOv1 target tensor + per-cell mask concatenated as 6076 bytes/image. See planning/archive/yolo_final.md and preprocess_pets_mosaic.py for the on-disk format. Only valid with lossKind := .yolov1Masked (or useYolov1 := true).

• brats : DatasetKind

  Brain-tumour segmentation on the Medical Segmentation Decathlon Task01_BrainTumour volumes (BraTS-derived). 2D axial slices: images are 240×240×4 (FLAIR / T1w / T1gd / T2w modalities as channels, z-scored per volume over brain voxels — no ImageNet normalization), labels are 240×240 uint8 per-pixel classes (0=background, 1=edema, 2=non-enhancing tumour, 3=enhancing tumour). Segmentation kind: labelBytesPerRecord = 240*240 selects .perPixelCE automatically. See preprocess_brats.py for the on-disk format and planning/archive/brats_demo.md for the demo plan.

• brats224 : DatasetKind

  The same BraTS data at 224×224, produced by preprocess_brats.py --size 224 (a center crop, not a resize — see fit_plane). Identical in every other respect: same 4 modalities, same patient split at seed 0, same slice selection, same 14,415/2,569 counts.

  224 exists for backbones with a /32 total stride, which 240 cannot serve (240/32 = 7.5). It is also ResNet-34's native ImageNet resolution. The crop is lossless for this dataset: across 4,000 sampled slices, the 8-pixel border it removes contains zero tumour voxels and zero brain voxels — MSD's volumes are skull-stripped and centered, so the margin is pure background. Used by demos/MainUnetBratsR34.lean.

@[implicit_reducible]

instance instReprDatasetKind : Repr DatasetKind

Equations

• instReprDatasetKind = { reprPrec := instReprDatasetKind.repr }

def instReprDatasetKind.repr : DatasetKind → Nat → Std.Format

Equations

• instReprDatasetKind.repr DatasetKind.mnist prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.mnist")).group prec✝
• instReprDatasetKind.repr DatasetKind.cifar10 prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.cifar10")).group prec✝
• instReprDatasetKind.repr DatasetKind.imagenette prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.imagenette")).group prec✝
• instReprDatasetKind.repr DatasetKind.pets prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.pets")).group prec✝
• instReprDatasetKind.repr DatasetKind.imagenet prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.imagenet")).group prec✝
• instReprDatasetKind.repr DatasetKind.petsDet prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.petsDet")).group prec✝
• instReprDatasetKind.repr DatasetKind.brats prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.brats")).group prec✝
• instReprDatasetKind.repr DatasetKind.brats224 prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "DatasetKind.brats224")).group prec✝

@[implicit_reducible]

instance instBEqDatasetKind : BEq DatasetKind

Equations

• instBEqDatasetKind = { beq := instBEqDatasetKind.beq }

def instBEqDatasetKind.beq : DatasetKind → DatasetKind → Bool

Equations

• instBEqDatasetKind.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

def ireeCompileArgs (mlirPath outPath : String) : IO (Array String)

IREE compile flags from environment. Defaults to CUDA (sm_86). Set IREE_BACKEND=rocm and IREE_CHIP=gfx1100 for AMD GPUs. Set IREE_BACKEND=llvm-cpu for CPU fallback (no chip needed).

Equations

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

def compileCheckB (name body : String) (stderrTake : Nat := 3000) : IO Bool

Compile one MLIR module and report. Writes body to .lake/build/{name}.mlir, runs iree-compile on it, prints the verdict, and returns true on success.

Lifted here from thirteen gates that each carried a byte-identical private copy. They differed only in how much of iree-compile's stderr they echoed (2000 / 3000 chars) — stderrTake keeps that adjustable, but no caller overrides it any more.

Equations

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

def compileCheck (name body : String) (stderrTake : Nat := 3000) : IO Unit

compileCheckB with the verdict discarded — the shape most gates want.

Equations

• compileCheck name body stderrTake = discard (compileCheckB name body stderrTake)

def tryCompile (src dst label : String) (stderrTake : Nat := 3000) : IO Unit

Compile src → dst, tolerating an absent compiler. Unlike compileCheck this takes an MLIR file that already exists and never throws: a missing iree-compile is reported and stepped over, so a smoke gate still runs on a machine without IREE installed.

Equations

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