Shared driver for the *-verified trainers

Every Main*Verified.lean trains a network on pre-rendered, audited StableHLO (verified_mlir/<slug>_{train_step,fwd}.mlir, emitted offline by tests/Test* from the proof stack) through the IREE FFI. Unlike the reference NetSpec/Train.lean path — which generates the MLIR at runtime — the verified path consumes a fixed codegen artifact, so a verified "model definition" is just:

• slug — which verified_mlir/*.mlir + which m.* functions to invoke,
• specs — the param layout ((dims, initKind), = the matching XLayout.specs),
• d0 — per-example input width, and
• data — which dataset/loader to feed it.

The architecture itself lives in the renderer + the audited VJP theorems; it is deliberately NOT re-expressed here. This file factors the ~100 lines of identical boilerplate (compile → sessions → load → init → train/eval loop) that every trainer used to copy. A trainer is now a VerifiedNet value + a VerifiedConfig + a one-line main, mirroring the shape of MainResnetTrain.lean.

NB the learning rate is baked into the rendered train-step MLIR — VerifiedConfig.lr is for the banner only; changing it does not change training (re-render to change lr).

inductive VerifiedData : Type

Which dataset a verified trainer runs on. Picks the loader, the eval-split name, and whether the training images need a 256²→224² center-crop per batch.

• mnist : VerifiedData

  MNIST idx files directly under dataDir (28×28×1, no crop).

• cifar : VerifiedData

  CIFAR-10 .bin records under dataDir/cifar-10 (32×32×3, no crop).

• imagenette : VerifiedData

  Imagenette under dataDir/imagenette — train stored at 256² (center-cropped to 224² per batch), val at 224².

• imagenet : VerifiedData

  Full 1000-class ImageNet, streamed from the generated tfds shim (handoff §2k). 1,281,167 train / 50,000 val at 224².

  Unlike every case above, this one is NOT preloaded: at f32 the train split is ~938 GiB of host RAM against a 188 GB box, so it cannot be. Batches arrive over a pipe from that net's own jax/.lake/build/generated_*_imagenet_shim.py (VerifiedNet.shimScript), already augmented, mean/std-normalized and flattened to (B, 3·224·224) — so the Lean side does no augmentation at all for this dataset, which is the point: there is exactly one definition of the transform and it is the one the JAX reference trainer uses.

  ⚠ Per net, and that is the part that was wrong until 2026-08-02: the script was hardcoded to ResNet-34's, so every net got RRC+hflip regardless of what its reference asked for. The transform is still single-definition — it is generated from the same TrainConfig the reference trainer runs — but WHICH definition is now a property of the net.

  The VAL split IS preloaded (49,920 imgs after tfds drop_remainder ⇒ 30 GB, which fits), so the eval loop is unchanged. 49,920 is the same count the reference run reported.

def instBEqVerifiedData.beq : VerifiedData → VerifiedData → Bool

Equations

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

@[implicit_reducible]

instance instBEqVerifiedData : BEq VerifiedData

Equations

• instBEqVerifiedData = { beq := instBEqVerifiedData.beq }

def instReprVerifiedData.repr : VerifiedData → Nat → Std.Format

Equations

• instReprVerifiedData.repr VerifiedData.mnist prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VerifiedData.mnist")).group prec✝
• instReprVerifiedData.repr VerifiedData.cifar prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VerifiedData.cifar")).group prec✝
• instReprVerifiedData.repr VerifiedData.imagenette prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VerifiedData.imagenette")).group prec✝
• instReprVerifiedData.repr VerifiedData.imagenet prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VerifiedData.imagenet")).group prec✝

@[implicit_reducible]

instance instReprVerifiedData : Repr VerifiedData

Equations

• instReprVerifiedData = { reprPrec := instReprVerifiedData.repr }

structure VerifiedNet : Type

A verified trainer: a pinned codegen artifact (slug) + its param layout (specs, d0, nClasses) + the dataset to run it on. See the module docstring.

• name : String

  Display name, e.g. "ResNet-34".

• mlirDir : String

  ⚠ Which directory this net's artifacts live in. Default verified_mlir/ — the CERTIFIED corpus, whose contents are pinned by scripts/regen_verified_mlir.sh check to exactly the set with a literal IO.FS.writeFile "verified_mlir/…" writer in Proofs/Codegen/.

  ⚠⚠ The width/batch SWEEP nets set .lake/build instead, and that is what keeps the pin possible. mlpG, cnnG and cifar8BnG render their artifact at run time from argv and immediately train on it — so those files are BUILD PRODUCTS, not committed renders. They used to be written into verified_mlir/ and 74 of them had been checked in: never loaded by anything, regenerated on every invocation, and invisible to the writer audit because that audit greps for a LITERAL path and these writers interpolate a slug. A directory that mixes a certified corpus with transients cannot be audited as either.

  ⚠ It is a field rather than a global because the read sites are per-net and there are 30 of them; one spelling in VerifiedNet beats 30 in the driver.

• slug : String

  Codegen slug: drives <mlirDir>/<slug>_{train_step,fwd}.mlir, .lake/build/<slug>_{ts,fwd}_v.vmfb, and the m.<slug>_{train_step,fwd} funcs.

• specs : Array (Array Nat × Nat)

  (dims, initKind) per param, in func-arg order — the matching XLayout.specs. initKind: 0 = He(fan-in), 1 = ones (γ), 2 = zeros (β / bias).

• d0 : Nat

  Per-example flattened input width (e.g. 3 * 224 * 224).

• nClasses : Nat

  Number of output classes.

• data : VerifiedData

  Dataset / loader selector.

• blurb : String

  One-line intro printed at startup (the prose banner). Carries the literal %LOWERER% where the transport belongs; print it with printBlurb, never with IO.println directly, so the banner names the lowerer that actually ran.

• lossSlot : Bool

  Does <slug>_train_step.mlir return the trailing report-only %loss scalar?

  A per-RENDER fact, not a driver-wide one. VerifiedNet.train used to append the slot unconditionally, which was true only of mlp and cnn (the two re-rendered for the chapter-2/3 loss carve-out) and wrong for every other net on this driver — resnet34, cifar8, cifar8_bn, cifar, cifar_bn, mobilenetv2, efficientnet, convnext, vit. Those all return parameters only, so the driver offered one destination too many and the G4 arity gate refused to run them.

  A wrong value here cannot corrupt anything: G4 compares the module's real output count against what the driver supplies and refuses on any mismatch.

• bnChannels : Array Nat

  Per-BN-layer channel counts, in forward order (empty for LayerNorm / no-BN nets). When non-empty, trainAdamSched threads running BN stats: the adam train step carries per-layer batch mean/var out in passthrough slots, the driver EMAs them, and eval uses <slug>_fwd_eval.mlir (affine BN with the running stats) instead of <slug>_fwd.mlir.

• dropKeeps : Array Float

  Stochastic-depth keep probabilities, one per drop site, in the render's signature order (planning/archive/stochastic_depth.md). Empty = the net has no drop sites, which is every net today except EfficientNet's *sd variants.

  ⚠ THE DRIVER OWNS THE RAMP, and that is deliberate rather than a shortcut: the emitted op is a pure per-example multiply and 1/keep_i is folded into the value supplied here, because a BAKED 1/keep_i and "the forward emits the sites too" cannot both hold (a ones scale would then compute x/keep_i, and the reference returns the branch untouched at eval). It is the same place %lr lives, for the same reason — one graph, many schedules.

  ⚠ It is therefore a SECOND hand-list against the renderer's enetDropIdxs, exactly like toSpecs == XLayout.specs. tests/TestDropPathRamp.lean is the #guard that pins the two; VerifiedSpec sits downstream of this file, so the renderer cannot share the definition by import without inverting the dependency.

• dropoutKeep : Option (Float × Nat)

  ▶ CLASSIFIER DROPOUT (recipe_gaps.md gap C) — (keep_prob, per-example width), or none when the net has none. EfficientNet-B0: (0.8, 1280) for the reference's dropout := 0.2 (jax/MainEfficientNetImagenet.lean:68).

  ⚠⚠ THE WIDTH IS HERE BECAUSE THE MASK IS PER-ELEMENT, WHICH IS THE WHOLE DIFFERENCE FROM dropKeeps. Stochastic depth's masks are tensor<Bxf32> — one value per example, so the driver needs no width at all. Dropout's is tensor<B×w×f32>, drawn per (example, feature), because the reference draws bernoulli(key, keep, x.shape) rather than the (B, 1, …, 1) shape. Every downstream difference — the blob shape, the draw count, the shard split — falls out of that one number, which is why it is carried rather than assumed to be net.d0 or nClasses. It is the CLASSIFIER'S INPUT width (EfficientNet's head channels), independent of the class count, so the Imagenette and ImageNet renders take the same mask shape.

  ⚠ keep_prob, not the drop rate: 1/keep is folded into the supplied mask by the driver, so the graph bakes no constant and the ones-mask forward is the exact identity (Proofs.dropout_ones_id). Same convention as dropKeeps, for the same reason.

• shimScript : String

  The generated ImageNet batch shim this net streams, as a bare filename under jax/.lake/build/ — e.g. "generated_vit_tiny_imagenet_shim.py". Required on every .imagenet net; ignored (and empty) on every other dataset, which loads from disk.

  ⚠⚠ THIS FIELD EXISTS BECAUSE ITS DEFAULT USED TO BE R34's, FOR EVERY NET. spawnShim hardcoded generated_resnet34_imagenet_shim.py and $SHIM_SCRIPT was set nowhere, so a "verified EfficientNet / ViT / ConvNeXt ImageNet run" streamed ResNet-34's augmentation — RandomResizedCrop + hflip and nothing else. Their references do not: EfficientNet's sets useAutoAugment, ViT's sets RandAugment m9/mstd0.5/inc1 + random erasing + repeated aug ×3, ConvNeXt's sets RandAugment + random erasing. The capability was there all along (JaxCodegen.generateShim honours every one of those flags); what was missing was the wiring, so the recipe matrix read ✅ on a capability rather than on the state.

  There is deliberately no fallback. An empty value on an .imagenet net REFUSES at spawn rather than substituting anything, because the failure this replaces was silent: the wrong augmentation compiles, streams, trains and descends. scripts/gen_shims.sh writes all five; $SHIM_SCRIPT still overrides with an explicit path, for a hand-placed or probe shim.

structure VerifiedConfig : Type

Training hyperparameters — the TrainConfig of the verified path. Mirrors the reference TrainConfig; kept as its own object so a net is a (spec, config) pair.

• epochs : Nat

  Number of training epochs.

• batchSize : Nat

  Minibatch size (a free runtime param — the MLIR's batch dim is dynamic).

• lr : Float

  Learning rate. DISPLAY ONLY — baked into <slug>_train_step.mlir; changing it here does not change training (re-render the MLIR to change lr).

• vitInit : Bool

  timm/DeiT ViT weight init — the verified peer of TrainConfig.vitInit, i.e. of the deit-init recipe the phase-2 reference run used (blueprint §9.6). Every weight at σ = 0.02 except the patch-embed conv on PyTorch's U(±1/√fan_in). Unlike lr, this is NOT display-only: init is host-side, so the flag genuinely changes training and no re-render is needed. Off by default — every other net keeps its seed reproducibility.

• bnMomentum : Float

  BatchNorm running-statistic decay — the verified peer of TrainConfig.bnMomentum, and the same TF sense: the weight on the OLD estimate, so timm's PyTorch momentum = 0.1 is 0.9 here. That field's docstring carries the per-net table and the timm audit; keep the two in step, since a phase-2 ↔ phase-4 gap here is invisible in every loss curve.

  Like vitInit and unlike lr, this is NOT display-only and needs no re-render: the graph emits raw per-layer BATCH stats and the host EMAs them (F32.ema below), so the decay never enters the MLIR. Under gradient accumulation the driver compensates to bnMomentum^(1/k) per micro-batch, matching the reference's generated _bn.

def wilson95 (correct nEval : Nat) : String

The 95% Wilson score interval on an accuracy, as lo–hi in percentage points.

⭐ Added 2026-08-30 because the eval line printed 3339/3925 = 85.070064% — six significant figures on a quantity a 3,925-image validation set resolves to about one. At p ≈ 0.85 and n = 3925 the 95% half-width is ±1.11 pt, so all but the leading three digits were decoration, and the ~1.3 pt epoch-to-epoch swings that get read as "the model moved" are inside it.

⚠ Wilson, not the normal approximation p ± z·√(p(1−p)/n). The normal form collapses to ±0 at correct = 0 or correct = n — it would print a perfectly certain 0.000000% on a run that scored nothing, which is exactly the LEAN_MLIR_SKIP_EVAL failure the line below this one already had to be taught to say out loud. Wilson stays finite there.

⚠ It is the MEASUREMENT error only — how well 3,925 images pin this model's accuracy. It says nothing about seed-to-seed training variance, which needs n runs, and it is the WRONG test for comparing two models scored on the SAME set: that comparison is paired, so it wants McNemar over LEAN_MLIR_DUMP_CORRECT's per-example bitmaps, which is far more powerful.

Equations

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

def VerifiedConfig.bnEmaWeight (cfg : VerifiedConfig) : Option Nat → Float

The weight F32.ema puts on the NEW batch, i.e. 1 − decay, given the accumulation factor: some k on an acc variant (the EMA fires once per MICRO-batch, so the decay is k-th-rooted for k chained updates to compose to one bnMomentum/optimizer-step update), none otherwise.

⭐ ONE definition, because it has three consumers — trainAdamSched's loop, its startup banner, and the fp8 trainer — and this file's own VerifiedVariant docstring is the record of what happens when a driver and its gate each keep a copy: an edit to the real expression cannot turn the copy red.

⚠ The == 0.99 arm returns the historic 0.01 DOUBLE rather than 1.0 - 0.99 (= 0.010000000000000009 — different bits, 9e-16 relative), so every net that leaves bnMomentum at its default is bit-identical across the knob's introduction. ⚠ some 0 is reachable — VerifiedVariant.accK returns 0 when it cannot parse a k out of the variant name — and lands on 1/0 = inf, pow → 0, weight 1.0, i.e. the running stats become the latest batch. That is the pre-knob behaviour, preserved deliberately rather than quietly repaired: a variant name whose k does not parse is already training at a wrong effective LR (see accK), and this should not be the thing that hides it.

Equations

• cfg.bnEmaWeight (some k) = 1.0 - cfg.bnMomentum.pow (1.0 / k.toFloat)
• cfg.bnEmaWeight none = if (cfg.bnMomentum == 0.99) = true then 1e-2 else 1.0 - cfg.bnMomentum

def VerifiedNet.paramShapes (n : VerifiedNet) : Array (Array Nat)

Param shapes in func-arg order (= specs dims).

Equations

• n.paramShapes = Array.map (fun (x : Array Nat × Nat) => x.fst) n.specs

def VerifiedNet.shapesBA (n : VerifiedNet) : ByteArray

Packed shape descriptors for the FFI (see packShapes).

Equations

• n.shapesBA = packShapes n.paramShapes

def VerifiedNet.nParams (n : VerifiedNet) : Nat

Total float count across all params.

Equations

• n.nParams = Array.foldl (fun (x1 x2 : Nat) => x1 + x2) 0 (Array.map (fun (s : Array Nat × Nat) => Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 s.fst) n.specs)

def VerifiedNet.xShape (n : VerifiedNet) (batch : Nat) : ByteArray

Packed x input shape [batch, d0].

Equations

• n.xShape batch = packXShape #[batch, n.d0]

LEAN_MLIR_VARIANT's axis predicates — ONE definition each

variant encodes five independent axes and every consumer recovers each with a string test on the name. tests/TestVariantPredicates.lean is the table of what each must read, and its docstring is the history: the naming has collided three times, each time between a PAIR of markers meeting rather than between a new marker and an old one.

⚠⚠ THEY LIVE HERE BECAUSE THE TEST USED TO PIN COPIES. trainAdamSched computed all five inline and TestVariantPredicates declared its own private def of each, so the table gated a transcription of the driver rather than the driver: an edit to the real predicate could not turn that file red. That is next_session_verified_trainer_code.md §5's lesson one level up — a gate on a definition is not a gate on the definition — and scoreCheckpoint needing the same region arithmetic is what made a third copy the alternative.

▶ The && !net.dropKeeps.isEmpty / && net.dropoutKeep.isSome conjuncts stay at the call sites: those are facts about the NET, not about the name, and folding them in here would make the predicate untestable from a string alone.

def VerifiedVariant.emaOn (v : String) : Bool

EMA shadow — a FOURTH [θ|m|v|ema] blob region, 5 scalars not 3.

Equations

• VerifiedVariant.emaOn v = v.startsWith "ema"

def VerifiedVariant.rmsOn (v : String) : Bool

RMSProp — the mean-square slot initialises to 1.0, not 0. ⚠ SUBSTRING, not prefix: the RMSProp+EMA spelling is emarms, which does not start with "rms" (planning/archive/ema.md's defect).

Equations

• VerifiedVariant.rmsOn v = decide ((v.splitOn "rms").length > 1)

def VerifiedVariant.sdOn (v : String) : Bool

Stochastic depth — N extra tensor<Bxf32> scale inputs. ⚠ The marker is drop and not sd because rms ++ dp spells rmsdp, which contains "sd" (planning/archive/stochastic_depth.md's defect).

Equations

• VerifiedVariant.sdOn v = decide ((v.splitOn "drop").length > 1)

def VerifiedVariant.cdOn (v : String) : Bool

Classifier dropout — ONE extra tensor<B×wxf32> mask input. ⚠ The marker is do and not dropout because dropout contains drop, so a dropout-only variant would read as a stochastic-depth one (recipe_gaps.md gap C).

Equations

• VerifiedVariant.cdOn v = decide ((v.splitOn "do").length > 1)

def VerifiedVariant.accOn (v : String) : Bool

Gradient accumulation — a FOURTH [θ|m|v|G] region, 5 scalars. ⚠⚠ SUBSTRING, not prefix: RSB-A3's composed optimizer is lambaccdp8x64bce, where lamb ++ acc puts the marker in the MIDDLE.

Equations

• VerifiedVariant.accOn v = decide ((v.splitOn "acc").length > 1)

def VerifiedVariant.lambOn (v : String) : Bool

LAMB — the per-tensor trust ratio (R34Opt.lambAccum), RSB-A3's optimizer. ⚠ SUBSTRING, not prefix, for accOn's reason one spelling over: the EMA form is emalamb…, which does not start with "lamb".

Equations

• VerifiedVariant.lambOn v = decide ((v.splitOn "lamb").length > 1)

def VerifiedVariant.bceOn (v : String) : Bool

BCE-with-logits (timm BinaryCrossEntropy), RSB's loss — not softmax CE. ⚠ SUBSTRING: the marker TRAILS the shape and is itself often trailed, by wd001 or bf16 (lambaccdp8x64wxclipbcebf16), so neither a prefix nor a suffix test finds it.

Equations

• VerifiedVariant.bceOn v = decide ((v.splitOn "bce").length > 1)

def VerifiedVariant.accK (v : String) : Nat

k, read back out of the name. The graph has 1/k BAKED in and the driver decides the apply cadence; a disagreement does not fail, it trains at a silently wrong effective learning rate. Parsed from AFTER the marker, not from a fixed offset — see accOn.

Equations

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

def VerifiedVariant.nRegions (v : String) : Nat

Blob regions: [θ|m|v], plus G (gradient accumulation) and/or E (the EMA shadow) — so 3, 4 or 5, and the two extras are INDEPENDENT. ⚠ A 3-region file loaded by a 4-region driver (or the reverse) misaligns EVERY parameter, so every consumer of a checkpoint sizes off this rather than off a literal.

⭐⭐ THIS USED TO BE if emaOn || accOn then 4 else 3, and the two features were mutually exclusive because of it — trainAdamSched threw on the pairing, and RSB-A2/A1 could not be rendered faithfully (their recipe sets BOTH gradAccumSteps := 4 and useEMA := true; planning/archive/verified_side_quest_counterparts.md §4a). The fifth region is what lifts that.

⚠⚠ G COMES BEFORE E, and that ordering is not free: at acc alone G is region 3 and at ema alone E is region 3, so every checkpoint written before this change still reads at the index it was written at. The reverse order would have silently re-homed every committed ema* blob.

Equations

• VerifiedVariant.nRegions v = (3 + if VerifiedVariant.accOn v = true then 1 else 0) + if VerifiedVariant.emaOn v = true then 1 else 0

def VerifiedVariant.nScalars (v : String) : Nat

Rank-0 scalar slots in the blob tail: lr,bc₁,bc₂, then %aup,%akeep (accumulation) and then %emad,%oemad (EMA) — so 3, 5 or 7, in that order. ⚠ Same independence and same ordering rule as nRegions: the accumulation pair keeps slots 3–4 and the EMA pair moves to 5–6 only when both are on, so neither single-axis layout moves.

Equations

• VerifiedVariant.nScalars v = (3 + if VerifiedVariant.accOn v = true then 2 else 0) + if VerifiedVariant.emaOn v = true then 2 else 0

def VerifiedVariant.emaRegion (v : String) : Option Nat

The blob index of the EMA shadow region, or none when the variant has no shadow. ⚠ It is not the literal 3 any more: under accumulation G takes region 3 and the shadow is region 4. scoreCheckpoint and the per-epoch eval both slice θ out of the blob with this, and a stale literal there does not fail — it scores the gradient accumulator as if it were weights and prints a plausible percentage off it.

Equations

• VerifiedVariant.emaRegion v = if VerifiedVariant.emaOn v = true then some (3 + if VerifiedVariant.accOn v = true then 1 else 0) else none

def VerifiedVariant.emaScalarOff (v : String) : Nat

Offset of the %emad,%oemad pair inside the scalar tail — 3 alone, 5 behind %aup,%akeep.

Equations

• VerifiedVariant.emaScalarOff v = 3 + if VerifiedVariant.accOn v = true then 2 else 0

def mkSession (mlirPath : String) : IO LowererSession

Open a session for one Lean-emitted graph, on whichever backend this binary dlopened (planning/archive/xla_pjrt_ladder.md).

• XLA — hand the .mlir straight to PJRT, which compiles it in-process. Nothing is written to disk.
• IREE — iree-compile the .mlir to a cache file first, then load that.

The cache path is derived from mlirPath rather than passed in. All 58 call sites used to supply one and every one of them computed the same thing from the same slug, so the argument was a second place for the name to be wrong and no place for it to be right. Deriving it also makes collisions impossible: two graphs cannot land on one cache file, which is the failure the target scoping below exists to prevent from the other direction.

Both backends consume the same verified_mlir/*.mlir — the emitter, the spec, and the §1a ties are identical. Only the trusted lowerer differs.

Equations

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

def mkParam (seed : Nat) (dims : Array Nat) (kind : Nat) (vitInit : Bool := false) (biasSigma : Option Float := none) (heFanIn : Bool := false) : IO ByteArray

Init one parameter from its (dims, initKind) spec, matching the JAX reference's initialisers — they are the oracle these nets are paired against:

• rank-4 conv kernel [oc, ic, kH, kW] → He fan-OUT, variance 2/(oc·kH·kW)
• rank-2 dense matrix [in, out] → Glorot, variance 2/(in + out)
• γ = 1 (kind 1), β / bias = 0 (kind 2)

⚠ Both weight cases CHANGED 2026-08-04. This used variance 2/fan_in for BOTH, where jax/Jax/Codegen.lean emits uniform(±√(6/fan_out)) for convs (variance 2/fan_out — torchvision's kaiming_normal_(mode='fan_out', nonlinearity='relu') convention for ResNet, emitConvBnInit) and uniform(±√(6/(fan_in+fan_out))) for dense (Glorot, emitDenseInit). The two paths had therefore never agreed on init, on any net. It is identical wherever ic == oc; the gaps are the stem (R34: fan_in 147 vs fan_out 3136 — 4.6× in σ), every stage-entry conv and 1×1 projection (2×), and every classifier (R34: 2/512 vs 2/1512, 1.7× in σ).

⚠⚠ This moves the init of every verified net, so no previously recorded accuracy is reproducible from its seed any more. It changes no committed artifact — init is host-side and no verified_mlir/ file mentions it.

⚠ The DISTRIBUTION still differs and is left alone deliberately: F32.heInit sums three uniforms (Bates-3, ≈ normal) where JAX draws one uniform. Variance is matched; shape is not. torchvision itself uses a normal here, so neither side is canonical on that axis, and changing the sampler would move every net for a second-order reason.

Equations

• One or more equations did not get rendered due to their size.
• mkParam seed dims 1 vitInit biasSigma heFanIn = F32.const (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize 1.0
• mkParam seed dims 2 vitInit none heFanIn = F32.const (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize 0.0
• mkParam seed dims 2 vitInit (some s) heFanIn = F32.heInit seed.toUSize (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize s

The ImageNet batch shim (handoff §2k)

Reads batches from JaxCodegen.generateShim's stdout. The shim owns the whole transform; this side only frames bytes, which is why there is no augmentation code here.

def readExact (h : IO.FS.Handle) (n : Nat) : IO ByteArray

Read EXACTLY n bytes, looping until they arrive. A pipe read returns what is available, not what was asked for — at 154 MB per batch a short read is the normal case, not the edge case, and treating one read as a batch silently misaligns the stream from then on.

Equations

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

def resolveShimScript (shimScript : String) : IO (Option System.FilePath)

Resolve a net's generated shim to a path on disk, or none. The candidate list is spawnShim's, factored out so the two callers cannot disagree about WHICH file they are reading — one of them decides the wire, the other spawns it.

Equations

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

def shimMixDefault (shimScript : String) : IO String

⭐ The SHIM_MIX default a generated shim BAKES — read out of the producer rather than restated here. "" when the shim cannot be found or declares nothing.

⚠ This is the mixup-λ lesson one layer up (§0.4 finding 3): recover a constant by READING it, not by fitting or re-declaring it. The alternative was a useMixup field on VerifiedNet duplicating what generateShim already baked from the same config — a second definition of one fact, which is the failure this repo keeps paying for. The shim text is the single source.

Equations

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

def spawnShim (shimScript split : String) (batch flat seed : Nat) (shard : Option (Nat × Nat) := none) (nclasses : Nat := 0) : IO IO.FS.Handle

Spawn the shim for one split and consume its preamble.

The preamble (LMSH | version | batch | flat) is checked rather than skipped: a batch or resolution mismatch between the render and the shim would otherwise read as garbage pixels and look like a broken net. Same reasoning as the FFI's G4 arity guard.

Equations

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

def readShimBatch (h : IO.FS.Handle) (batch flat : Nat) (nclasses : Nat := 0) : IO (ByteArray × ByteArray)

One batch off the wire: int32[batch] labels then float32[batch*flat] images, in that order (the shim writes labels first so a partial record is detectable at the smaller read).

Equations

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

def readUpTo (h : IO.FS.Handle) (n : Nat) : IO ByteArray

Read up to n bytes, returning what actually arrived instead of throwing at EOF. The peer of readExact, and the only difference is which of "short read" and "clean end of stream" it treats as the error.

Equations

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

def readShimBatchPartial (h : IO.FS.Handle) (batch flat : Nat) (nclasses : Nat := 0) : IO (ByteArray × ByteArray × Nat)

One shim batch, tolerating a SHORT FINAL BATCH — the validation-split reader.

⚠⚠ Why this exists (2026-08-14). The val pipeline used drop_remainder=True, so ImageNet's 50,000 images batched at 256 gave 195 full batches and 80 images were thrown away. Every top-1 this repo has quoted for an ImageNet net is therefore over 49,920, where timm's validate.py scores all 50,000 — a difference of 0.16% that is not an error bar, it is a different denominator. The shim now sets drop_remainder=training, which puts a partial batch on the wire that readExact refuses by construction ("shim closed the pipe after N of M bytes"). This reader accepts it.

Returns (img, lbl, rows) where rows ≤ batch, and rows = 0 means the stream ended cleanly.

⚠ It reads LABELS FIRST, matching the wire order, and infers rows from the label read — the label record is 4 bytes (or 4·nclasses) against the image's 4·flat, so a truncated stream is far more likely to be caught mid-image than mid-label. Inferring from the SMALLER record and then demanding exactly that many image bytes turns a torn write into a loud failure instead of a silently short batch.

⭐ No MLIR changes. The eval graph keeps its baked batch width: F32.sliceImagesPad zero-pads the tail up to it and the eval loop scores min bs (nEval − bi·bs) real rows, so the pad never reaches the accuracy count. That is safe because eval normalises PER EXAMPLE everywhere — running-stat BN through @<slug>_fwd_eval, LayerNorm through @<slug>_fwd. ⚠ The one exception is LEAN_MLIR_EVAL_BATCHSTATS=1, which scores through @<slug>_fwd with BATCH statistics: there the zero rows WOULD shift the real rows' normalisation. That flag is a declared diagnostic, and the drain refuses to keep the tail under it.

Equations

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

def spawnShimSharded (shimScript split : String) (batch flat seed n : Nat) (nclasses : Nat := 0) : IO (Array IO.FS.Handle)

Spawn n shim processes over disjoint shards of one split, and read them round-robin.

Why this exists. One shim process tops out at ~1,530 img/s (measured 2026-08-01, bs128, marginal so TF startup is out of it). A 4-replica ViT step consumes 512 images in 264 ms, i.e. ~1,940 img/s, so a single producer would make the GPUs wait — the first config in this repo where the loader, not the device, is the ceiling (R34/ImageNet at bs256 needs only ~380). Measured aggregate: 2 processes 1.71×, 4 processes 2.36× on this 32-core box, so two clear the requirement with margin.

What it does NOT do is add a second definition of the transform. Each worker runs the same generated shim with SHIM_SHARD=i/n, which selects which examples it emits (ds.shard, before the map) and leaves _pp — the crop, flip and normalization — untouched. A hand-written loader here would be §2a's double-writer disease applied to the data path.

⚠ Round-robin over BATCHES is not the unsharded stream. ds.shard interleaves elements, so taking whole batches from each worker in turn gives a different batch composition than one producer would. Both are valid shuffled streams over the same epoch of data, and each worker shuffles its own slice with the pipeline's own seed — but the two are not byte-comparable, so a determinism hash from the unsharded config does not carry to a sharded one. Re-run SHIM_HASH per shard if you need that property.

⚠ Each worker gets a DISTINCT seed (seed + i). With one shared seed every worker draws the same augmentation sequence, and since the shards hold different images that is not a correctness bug — but it needlessly correlates the crops across workers.

Equations

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

def readShimBatchRR (hs : Array IO.FS.Handle) (k batch flat : Nat) (nclasses : Nat := 0) : IO (ByteArray × ByteArray)

Round-robin read: batch k comes from worker k % n. readExact already blocks until a whole record has arrived, so a slow worker throttles rather than corrupting — the framing cannot slip.

Equations

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

def VerifiedNet.ckptPathFor (net : VerifiedNet) (variant : String) : IO String

Where this (net, variant) writes and resumes its checkpoint.

Scoped by BACKEND: without the suffix an XLA run would happily resume from an IREE checkpoint and vice versa, silently fusing two trajectories into one while looking completely normal on screen (planning/archive/xla_pjrt_ladder.md §3). $LEAN_MLIR_CKPT_TAG appends a run-scoped suffix — without it every pass of the same (net, variant, backend) shares ONE path, so the parallel sweeps planning/archive/chapter_makeover.md §3c mandates cannot be run: concurrent passes clobber each other's blob, and a later pass resumes from an earlier one's finished epoch 40.

⚠ A function because scoreCheckpoint has to land on the SAME path the trainer wrote, and "score the checkpoint the run just finished" is that tool's zero-argument case. Two spellings of this string would not fail — they would score a file that is not there, or worse, an older one from a different tag.

Equations

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

def VerifiedNet.printBlurb (net : VerifiedNet) : IO Unit

Print the startup banner with %LOWERER% resolved to the lowerer that actually ran.

Every net's blurb used to hard-code its transport, so the banner was a claim about the build rather than about the run. Three of the seven print sites patched it at run time with a .replace "IREE FFI" "XLA/PJRT" and the other four printed it raw, which meant a net whose blurb still said IREE announced IREE while training on XLA. The placeholder moves the decision to one place, and every caller is correct by construction.

Equations

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

def VerifiedNet.train (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Train a VerifiedNet end-to-end on its proof-rendered StableHLO: compile both MLIRs → lowerer sessions → load data → He/spec init → SGD train + eval loop. The SGD update (and lr) are baked into <slug>_train_step.mlir; we only feed batches.

Equations

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

def VerifiedNet.trainAdamPacked (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

AdamW training driver — threads the first/second moment buffers as a single packed [θ|m|v] param blob through the generic FFI (n_params = 3k; the moments ride in the params slot, so the prebuilt .so is unchanged), against the baked-hyperparameter packed render @<slug>_adam_train_step (ViTRender.vitTrainStepModuleAdamPacked, optimizer = Proofs.adamWParam). Moments init to 0; eval reads the θ slice (first nParams floats). The Adam analogue of VerifiedNet.train.

Equations

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

def VerifiedNet.trainAdamSched (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (baseLR β1 β2 : Float) (warmupEpochs : Nat) (variant : String := "adam") (expDecayRate : Float := 0.0) (expDecayEpochs : Float := 1.0) (emaDecay : Float := 0.9999) : IO Unit

Scheduled AdamW driver (Phase 2) — trainAdamPacked with a runtime LR and bias correction. lr/bc₁/bc₂ ride as three rank-0 scalar params in the blob tail ([θ|m|v|lr|bc₁|bc₂], the FFI takes no scalar slot) and are returned unchanged; the host recomputes them each step: cosine decay + linear warmup for lr, and bc₁=1−β₁ᵗ, bc₂=1−β₂ᵗ (proper bias correction). Drives ViTRender.vitTrainStepModuleAdamSched.

expDecayRate > 0 selects the EfficientNet/MobileNetV2 exponential schedule over cosine: after warmup, lr = baseLR · rate^((epoch − warmupEpochs)/decayEpochs). Both references use it (mnv2 ×0.98 per epoch, EfficientNet ×0.97 every 2.4), and it is recipe_gaps.md Tier C — a driver item, not a render one, because lr is already a runtime operand. Default 0.0 keeps cosine, so every existing call site is unchanged.

RMSProp variants also need a different INITIAL STATE, which is the other half of that gap and is handled below — see rmsprop.

Equations

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

def VerifiedNet.scoreCheckpoint (net : VerifiedNet) (dataDir variant ckptPath : String) (region : String := "auto") : IO Unit

Score a checkpoint, standalone — the eval half of trainAdamSched with no training in front of it (planning/archive/next_session_verified_trainer_code.md §2).

Until this existed a verified accuracy could only be produced in training, and only for the weights that happened to be live at that moment. The JAX side has six eval_*_full50k.py; this is the verified peer, and it is a FACTORING job rather than new machinery — no new MLIR, no new ops, no renderer work. Every piece already existed inside the eval half:

need	reused
drain the val split	loadData (all 50,000 as of ccca380)
the eval graph	mkSession on <slug>_fwd_eval, or the _fwd chain for the LN nets
eval batch AND width	fwdRenderedShape, one parse of one declaration
batching a short tail	F32.sliceImagesPad + min evalBs (nEval − bi·evalBs)
forward	LowererSession.forwardF32
metrics	F32.argmaxN (top-1), F32.rankOf (top-5)

⭐ THE GATE IS AN EQUALITY, NOT A SMOKE TEST. For the same checkpoint at the same region, the number printed here must equal the one the training run printed for that epoch — same denominator, same batching, same graph. It is available today on ConvNeXt and ViT, which have nBnStats = 0 and therefore carry their whole eval state in the checkpoint.

⚠⚠ BN NETS ARE REFUSED, LOUDLY, AND THAT IS THE POINT. The checkpoint is exactly [θ|m|v(|ema)]; the BN running mean/var are NOT in it — they are "reset per process and rebuilt within an epoch" (see runningBnStats). In-training eval works because the statistics have been accumulating all epoch. A fresh process reading a .bin has ZEROS, and @<slug>_fwd_eval then normalises by them: not a slightly-off number, garbage that still prints as a plausible-looking percentage. So R50/R34/MNv2/EfficientNet/MNv4 throw here rather than score, until §2b lands the stats in the checkpoint (format) plus --recalibrate (the fallback, and the only one of the two that can reach A3's finished checkpoint).

⭐ region is what one checkpoint cannot otherwise yield: the driver picks live-or-shadow at TRAIN time (emaLiveBn), so an EMA run reports one of the two numbers and discards the other. timm reports the shadow and RSB-A2 sets emaDecay := 0.9999, so without this an A2 result is not quotable the way its reference is. "auto" = the shadow when the variant has one, matching what the training run would have scored; "live" and "ema" name it explicitly.

Equations

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

def VerifiedNet.trainLinear (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Train driver for the 2-parameter linear path (Chapter 1). The verified @<slug>_train_step takes W0/b0 as separate arguments (linearTrainStepV), weights are zero-initialized, and the loss/lr are baked into the MLIR — distinct from the packed-params, He-init train above. Only the linear classifier uses this; shares compileVmfb / loadData / the eval pass with the main driver.

Equations

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

def VerifiedNet.attackPgdMlp (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Phase-3 PGD attack on the verified MNIST MLP (planning/archive/robustness.md). Trains the 784→512→512→10 ReLU MLP on the proof-rendered SGD step, then attacks through IREE with the proven mlpInputGrad VJP kernel. The Lipschitz certificate is the product of the three layers' spectral norms — where the bound (and so the cert) goes loose.

Equations

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

def VerifiedNet.attackPgdSpectralMlp (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) : IO Unit

Spectral-norm-constrained training of the verified MNIST MLP (planning/archive/robustness_ladder.md, the research lever). Trains the 784→512→512→10 net with projected SGD onto the spectral ball — after every K proof-rendered steps (and once at the end) each weight Wᵢ is rescaled to ‖Wᵢ‖₂ ≤ c (projectSpectral) — then runs the same cert ≤ TRUE ≤ PGD sandwich. Sweeps a few caps c (plus an unconstrained baseline) so the table shows the trade: shrinking c pulls the global L = ∏‖Wᵢ‖₂ down (L ≤ c³), turning the vacuous product certificate non-vacuous — at the cost of clean accuracy. The empirical face of lipschitz_margin_certified_radius (LeanMlir/Proofs/Certificates/LipschitzCert.lean): smaller L ⇒ larger certified radius m/(√2·L). The verified CE gradient stays in the proven kernel; the projection is host-side weight rescaling only.

Equations

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

def VerifiedNet.attackPgdConvNet (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (genKernel : Nat → Float → Float → Bool → String) (withCert : Bool := true) : IO Unit

Generic conv-net PGD attack (planning/archive/robustness_ladder.md). Trains any packed conv net on its proof-rendered SGD step, then attacks through IREE with genKernel — the full proven backward (conv input-VJPs + maxpool select_and_scatter-backs, mirroring the net's <slug>_train_step.mlir) run to dx. Certificate = the conv-aware spectral-norm product (specNormConvTapSum for convs × specNormW for denses; ReLU/maxpool are 1-Lipschitz) — astronomically loose, the depth-cliff. genKernel and net.slug select the architecture (genCnnPgdStep/MNIST-CNN, genCifarPgdStep/CIFAR-CNN).

Equations

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

def VerifiedNet.attackPgdCnn (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

PGD attack on the verified MNIST CNN (the first conv rung).

Equations

• net.attackPgdCnn cfg dataDir = net.attackPgdConvNet cfg dataDir genCnnPgdStep✝

def VerifiedNet.attackPgdCifar (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

PGD attack on the verified CIFAR-10 CNN (the deeper conv rung: 4 conv + 2 pool + 3 dense).

Equations

• net.attackPgdCifar cfg dataDir = net.attackPgdConvNet cfg dataDir genCifarPgdStep✝

def VerifiedNet.attackPgdCifarBn (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

PGD attack on the verified CIFAR-10 CNN + per-channel (instance) BatchNorm (cifar_bn). The BN input-VJP rung — genCifarBnPgdStep runs the proven backward through 4 instance-norm layers (the BN grad-input 3-term formula). Certificate skipped (withCert := false): instance norm's Lipschitz is data-dependent (γ·istd), a separate problem from the conv-product.

Equations

• net.attackPgdCifarBn cfg dataDir = net.attackPgdConvNet cfg dataDir genCifarBnPgdStep✝ false

def VerifiedNet.attackPgdSpectralConvNet (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) (genKernel : Nat → Float → Float → Bool → String) : IO Unit

Spectral-norm-constrained training of the verified MNIST CNN (planning/archive/robustness_ladder.md, the gap-shrinking lever applied to the conv net). The CNN sibling of attackPgdSpectralMlp: projected SGD onto the spectral ball — every K proof-rendered steps (and once at the end) projectSpectral caps both the dense ‖Wᵢ‖₂ and the conv tap-sum bound at c — then the cert ≤ TRUE ≤ PGD sandwich (PGD via genKernel, cert = the conv-aware product). Harder than the MLP: it's a k-layer product (L ≤ cᵏ) and the conv tap-sum is a loose bound, so projection over-penalizes the convs — the cert needs a tighter c (and pays more clean accuracy) than the MLP did, and certifies only at smaller radii. The honest "depth + loose conv-norm ⇒ certifying the conv net is harder." Generic over genKernel/net.slug (MNIST-CNN, CIFAR-CNN).

Equations

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

def VerifiedNet.attackPgdSpectralCnn (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) : IO Unit

Spectral-norm-constrained training of the verified MNIST CNN.

Equations

• net.attackPgdSpectralCnn cfg dataDir caps = net.attackPgdSpectralConvNet cfg dataDir caps genCnnPgdStep✝

def VerifiedNet.attackPgdSpectralCifar (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) : IO Unit

Spectral-norm-constrained training of the verified CIFAR-10 CNN (7-layer product).

Equations

• net.attackPgdSpectralCifar cfg dataDir caps = net.attackPgdSpectralConvNet cfg dataDir caps genCifarPgdStep✝

Randomized-smoothing statistics (Cohen–Rosenfeld–Kolter 2019)

The pieces the smoothing certificate needs, in pure Float (no kernel, no Mathlib): the probit Φ⁻¹ for the radius σ·Φ⁻¹(p_A), and a sound Clopper–Pearson lower confidence bound on p_A (a genuine 1−α lower bound, not an approximation — a certificate must under- estimate). CP is built bottom-up from the regularized incomplete beta Iₓ(a,b).

def VerifiedNet.smoothCertify (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (sigmas : List Float) : IO Unit

Randomized-smoothing certificate (Cohen–Rosenfeld–Kolter 2019, planning/archive/robustness_ladder.md §3) — the depth-INDEPENDENT cert, and the answer where the Lipschitz product is hopeless. The smoothed classifier ĝ(x) = argmax_c P[f(x+η)=c], η ~ N(0,σ²I), is certified robust at L2 radius σ·Φ⁻¹(p_A) where p_A is a lower bound on the top class's noise probability. It's forward-only: no new kernel, no input-VJP — just sample n noisy copies, run the existing proof-rendered <slug>_fwd, count argmax votes, Clopper–Pearson lower-bound p_A. The base classifier is trained with matched Gaussian augmentation (every batch corrupted with N(0,σ²I) host-side before the proof-rendered SGD step — the forward/backward graph is untouched), the Cohen recipe. Architecture-agnostic + depth-independent, so it certifies a non-vacuous radius on the very nets (CIFAR, deep) where ∏‖Wᵢ‖₂ is astronomically loose. Generic over any VerifiedNet (fwd + train-step only).

n (SMOOTH_N, default 10000 — Cohen's large-n regime) is the estimation budget and the only honest tightening lever: the per-point radius is capped at σ·Φ⁻¹(α^(1/n)) (a unanimous vote still only certifies p_A ≤ α^(1/n)), so larger n lifts the ceiling and tightens the CP bound toward the true noise-probability — bigger certified radii at the same 1−α guarantee.

Equations

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

def VerifiedNet.attackPgd (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Phase-3 PGD adversarial attack on the verified linear classifier (planning/archive/robustness.md). Trains via the proof-rendered train step, then attacks through the real IREE pipeline: each PGD step's input gradient is computed by the genLinearPgdStep StableHLO kernel (the proven dx = (softmax−onehot)·Wᵀ VJP) on the GPU. Reports clean vs L∞-PGD adversarial accuracy over an eps sweep.

Equations

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

def VerifiedNet.trainLinearE4M3 (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

fp8 (E4M3) Lean trainer — the low-precision sibling of trainLinear.

Keeps fp32 master weights and, each step, projects the weights (per-output-column) and the activations (per-tensor) onto the E4M3 grid (LeanMlir/E4M3Quant.lean), runs the same verified @<slug>_train_step kernel (the matmul accumulates in fp32 — the dotMixed model: u_leaf = E4M3, u_acc = fp32), and applies the recovered gradient delta to the fp32 master via addDelta (master += Wout − Wq = master − lr·∇). The MLIR and FFI are unchanged: fp8 here is host-side operand byte-prep, exactly the §3b render-tie model (Proofs/E4M3Fold.lean). Eval runs the fp32 master through @<slug>_fwd (the "fp32-infer" accuracy of the fp8-trained model, mirroring scripts/mnist_e4m3_demo.py).

Run (GPU): IREE_BACKEND=rocm .lake/build/bin/mnist-linear-e4m3-verified data

Equations

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

def VerifiedNet.trainE4M3 (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

fp8 (E4M3) packed-params trainer — the low-precision sibling of VerifiedNet.train, for the depth>1 nets (MLP, CNN). Keeps fp32 master params and, each step, projects every weight slot onto the E4M3 grid (dense per-output-column, conv per-output-channel; biases kept fp32 — F32E4M3.quantPackedParams) and the input per-tensor, runs the same verified @<slug>_train_step (fp32 accumulate inside), and folds the gradient delta back into the master with addDelta over the whole packed buffer (master += out − paramsQ: weight slots get −lr·∇, bias slots the exact update). MLIR/FFI unchanged.

Scope (honest): host-side prep reaches weights + the input activation only. The intermediate activations (relu/pool/flatten outputs feeding the deeper matmuls) and the backward-chain cotangents are computed inside the fused kernel and stay fp32 — quantizing them needs in-graph E4M3 ops (the next, codegen-level step), not host byte-prep. So this is honest fp8 weights + fp8 input, fp32 intermediates. Eval runs the fp32 master.

Run (GPU): IREE_BACKEND=rocm .lake/build/bin/mnist-mlp-e4m3-verified data

Equations

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

def VerifiedNet.trainAdamSchedE4M3 (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (baseLR β1 β2 : Float) (warmupEpochs : Nat) (variant : String := "adam") (expDecayRate : Float := 0.0) : IO Unit

fp8 (E4M3) variant of trainAdamSched — runs the Adam / Nesterov-momentum optimizer demos in fp8. Keeps an fp32 master [θ|m|v]; each step projects the weight third θ onto the E4M3 grid (quantPackedParams: dense per-column, conv per-channel; biases fp32) and the input per-tensor, runs the same verified @<slug>_<variant>_train_step (the optimizer is baked into the MLIR, so fp8 needs no new module — operand byte-prep only; fp32 accumulate), and folds the optimizer-step delta back into the fp32 master θ (addDelta), keeping the returned m'/v' moments in fp32. Distinct _e4m3 checkpoint (won't resume an fp32 run); honors LEAN_MLIR_MAX_EPOCHS. Same scope as trainE4M3: fp8 weights + input, fp32 intermediates / moments.

Equations

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