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 the renderers in
LeanMlir/Proofs/Codegen/ and regenerated by scripts/regen_verified_mlir.sh) through the runtime FFI (PJRT by default, IREE optionally). 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— whichverified_mlir/*.mlir+ whichm.*functions to invoke,specs— the param layout ((dims, initKind), = the matchingXLayout.specs),d0— per-example input width, anddata— 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 identical
boilerplate (compile → sessions → load → init → train/eval loop) shared by every trainer.
A trainer is a VerifiedNet value + a VerifiedConfig + a one-line main, the same shape
as a NetSpec trainer.
This is the training driver alone (plus its fp8 E4M3 variants). The PGD attacks and spectral-norm
studies are Verified.Attack (on Verified.PgdGen's kernels), the smoothing certificate is
Verified.Smoothing. An entry point imports Verified.NetsCore (the specs) and whichever of
these three it runs.
NB VerifiedConfig carries no learning rate. The SGD-inline train steps bake it
into the rendered MLIR (re-render to change it); the Adam-family steps take it as a runtime operand,
from trainAdamSched's own baseLR argument and schedule.
References #
- Loshchilov & Hutter 2019, Decoupled Weight Decay Regularization (AdamW). https://arxiv.org/abs/1711.05101
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". - name : String
Display name, e.g.
"ResNet-34". - mlirDir : String
Which directory this net's artifacts live in. Default
verified_mlir/— the committed corpus, whose contentsscripts/regen_verified_mlir.sh checkpins to exactly the set with a literalIO.FS.writeFile "verified_mlir/…"writer in LeanMlir/Proofs/Codegen/.The width/batch sweep nets (
mlpG,cnnG,cifar8BnG) set.lake/buildinstead: they render their artifact at run time from argv and train on it at once, so those files are build products, not committed renders. Keeping them out ofverified_mlir/is what lets the writer audit pin that directory — the audit greps for a literal path, and these writers interpolate a slug.A field rather than a global because the read sites are per-net.
- mlirDir : String
Which directory this net's artifacts live in. Default
verified_mlir/— the committed corpus, whose contentsscripts/regen_verified_mlir.sh checkpins to exactly the set with a literalIO.FS.writeFile "verified_mlir/…"writer in LeanMlir/Proofs/Codegen/.The width/batch sweep nets (
mlpG,cnnG,cifar8BnG) set.lake/buildinstead: they render their artifact at run time from argv and train on it at once, so those files are build products, not committed renders. Keeping them out ofverified_mlir/is what lets the writer audit pin that directory — the audit greps for a literal path, and these writers interpolate a slug.A field rather than a global because the read sites are per-net.
- slug : String
Codegen slug: drives
<mlirDir>/<slug>_{train_step,fwd}.mlir,.lake/build/<slug>_{ts,fwd}_v.vmfb, and them.<slug>_{train_step,fwd}funcs. - slug : String
Codegen slug: drives
<mlirDir>/<slug>_{train_step,fwd}.mlir,.lake/build/<slug>_{ts,fwd}_v.vmfb, and them.<slug>_{train_step,fwd}funcs. (dims, initKind)per param, in func-arg order — the matchingXLayout.specs.initKind: 0 = random weight (conv: He fan-out; dense: Glorot; seemkParam), 1 = ones (γ), 2 = zeros (β / bias), 3 = 1e-6 (layer scale), 4 = zero-γ (0 underzeroGammaInit, else 1), 5 = embedding (σ 0.02 undervitInit, else 0), 7 = SE FC (var 2/out underseFanOutInit, else Glorot).- d0 : Nat
Per-example flattened input width (e.g.
3 * 224 * 224). - d0 : Nat
Per-example flattened input width (e.g.
3 * 224 * 224). - nClasses : Nat
Number of output classes.
- nClasses : Nat
Number of output classes.
- data : VerifiedData
Dataset / loader selector.
- 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 withprintBlurb, never withIO.printlndirectly, so the banner names the lowerer that actually ran. - blurb : String
One-line intro printed at startup (the prose banner). Carries the literal
%LOWERER%where the transport belongs; print it withprintBlurb, never withIO.printlndirectly, so the banner names the lowerer that actually ran. - lossSlot : Bool
Does
<slug>_train_step.mlirreturn the trailing report-only%lossscalar?A per-render fact, not a driver-wide one: true for the
mlpandcnnrenders, which emit the loss, and false for the nets whose train step returns parameters only.A wrong value here cannot corrupt anything: the G4 arity gate compares the module's real output count against what the driver supplies and refuses on any mismatch.
- lossSlot : Bool
Does
<slug>_train_step.mlirreturn the trailing report-only%lossscalar?A per-render fact, not a driver-wide one: true for the
mlpandcnnrenders, which emit the loss, and false for the nets whose train step returns parameters only.A wrong value here cannot corrupt anything: the G4 arity gate compares the module's real output count against what the driver supplies and refuses on any mismatch.
Per-BN-layer channel counts, in forward order (empty for LayerNorm / no-BN nets). When non-empty,
trainAdamSchedthreads 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.Per-BN-layer channel counts, in forward order (empty for LayerNorm / no-BN nets). When non-empty,
trainAdamSchedthreads 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.Stochastic-depth keep probabilities, one per drop site, in the render's signature order. Empty = no drop sites. Consumed only by
*drop*variants (VerifiedVariant.sdOn).The driver owns the ramp, deliberately: the emitted op is a pure per-example multiply and
1/keep_iis folded into the value supplied here, because a BAKED1/keep_iand "the forward emits the sites too" cannot both hold (a ones scale would then computex/keep_i, and the reference returns the branch untouched at eval). It is the same place%lrlives, for the same reason — one graph, many schedules.It is therefore a second hand-list against the renderer's
enetDropIdxs, liketoSpecs == XLayout.specs.tests/TestDropPathRamp.leanis the#guardthat pins the two;VerifiedSpecsits downstream of this file, so the renderer cannot share the definition by import without inverting the dependency.Stochastic-depth keep probabilities, one per drop site, in the render's signature order. Empty = no drop sites. Consumed only by
*drop*variants (VerifiedVariant.sdOn).The driver owns the ramp, deliberately: the emitted op is a pure per-example multiply and
1/keep_iis folded into the value supplied here, because a BAKED1/keep_iand "the forward emits the sites too" cannot both hold (a ones scale would then computex/keep_i, and the reference returns the branch untouched at eval). It is the same place%lrlives, for the same reason — one graph, many schedules.It is therefore a second hand-list against the renderer's
enetDropIdxs, liketoSpecs == XLayout.specs.tests/TestDropPathRamp.leanis the#guardthat pins the two;Verified.Specsits downstream of this file, so the renderer cannot share the definition by import without inverting the dependency.Classifier dropout —
(keep_prob, per-example width), ornonewhen the net has none. EfficientNet-B0:(0.8, 1280)for the reference'sdropout := 0.2(jax/MainEfficientNetImagenet.lean).The width is here because the mask is per-element, which is the difference from
dropKeeps. Stochastic depth's masks aretensor<Bxf32>— one value per example, so the driver needs no width at all. Dropout's istensor<B×w×f32>, drawn per (example, feature), because the reference drawsbernoulli(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 benet.d0ornClasses. 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/keepis 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 asdropKeeps, for the same reason.Classifier dropout —
(keep_prob, per-example width), ornonewhen the net has none. EfficientNet-B0:(0.8, 1280)for the reference'sdropout := 0.2(jax/MainEfficientNetImagenet.lean).The width is here because the mask is per-element, which is the difference from
dropKeeps. Stochastic depth's masks aretensor<Bxf32>— one value per example, so the driver needs no width at all. Dropout's istensor<B×w×f32>, drawn per (example, feature), because the reference drawsbernoulli(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 benet.d0ornClasses. 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/keepis 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 asdropKeeps, 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.imagenetnet; ignored (and empty) on every other dataset, which loads from disk.Each net's shim carries its own reference's augmentation (EfficientNet's sets
useAutoAugment, ViT's RandAugment m9/mstd0.5/inc1 + random erasing + repeated aug ×3, ConvNeXt's RandAugment + random erasing), so streaming another net's shim trains on the wrong recipe.There is deliberately no fallback: an empty value on an
.imagenetnet refuses at spawn, because a wrong shim fails silently — the wrong augmentation compiles, streams, trains and descends.scripts/gen_shims.shwrites all five;$SHIM_SCRIPTstill overrides with an explicit path, for a hand-placed or probe shim. - 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.imagenetnet; ignored (and empty) on every other dataset, which loads from disk.Each net's shim carries its own reference's augmentation (EfficientNet's sets
useAutoAugment, ViT's RandAugment m9/mstd0.5/inc1 + random erasing + repeated aug ×3, ConvNeXt's RandAugment + random erasing), so streaming another net's shim trains on the wrong recipe.There is deliberately no fallback: an empty value on an
.imagenetnet refuses at spawn, because a wrong shim fails silently — the wrong augmentation compiles, streams, trains and descends.scripts/gen_shims.shwrites all five;$SHIM_SCRIPTstill overrides with an explicit path, for a hand-placed or probe shim.
Instances For
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.
- epochs : Nat
Number of training epochs.
- batchSize : Nat
Minibatch size (a free runtime param — the MLIR's batch dim is dynamic).
- batchSize : Nat
Minibatch size (a free runtime param — the MLIR's batch dim is dynamic).
- valEveryEpochs : Nat
Validate every N epochs, plus always the last epoch the process runs — the verified peer of
TrainConfig.valEveryEpochs, which the reference's S/B/ViT-S/B/MobileNetV4 ImageNet configs set to 5. N ≤ 1 keeps every-epoch validation. ImageNet path (trainAdamSched) only, and only the eval pass: the checkpoint is still written every epoch.LEAN_MLIR_VAL_EVERY=<n>overrides it at launch, likeLEAN_MLIR_MAX_EPOCHS. - valEveryEpochs : Nat
Validate every N epochs, plus always the last epoch the process runs — the verified peer of
TrainConfig.valEveryEpochs, which the reference's S/B/ViT-S/B/MobileNetV4 ImageNet configs set to 5. N ≤ 1 keeps every-epoch validation. ImageNet path (trainAdamSched) only, and only the eval pass: the checkpoint is still written every epoch.LEAN_MLIR_VAL_EVERY=<n>overrides it at launch, likeLEAN_MLIR_MAX_EPOCHS. - vitInit : Bool
timm/DeiT ViT weight init — the verified peer of
TrainConfig.vitInit, i.e. of thedeit-initrecipe of the JAX reference. Every weight at σ = 0.02 except the patch-embed conv on PyTorch'sU(±1/√fan_in). Unlikelr, 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. - vitInit : Bool
timm/DeiT ViT weight init — the verified peer of
TrainConfig.vitInit, i.e. of thedeit-initrecipe of the JAX reference. Every weight at σ = 0.02 except the patch-embed conv on PyTorch'sU(±1/√fan_in); the CLS token and positional embedding (init kind 5) at σ = 0.02 too, where without the flag they are 0. Init is host-side, so the flag changes training with no re-render. Off by default — every other net keeps its seed reproducibility. - cnxInit : Bool
ConvNeXt
_init_weights— the verified peer ofTrainConfig.cnxInit.trunc_normal_(std=0.02)on every conv and the head; biases 0, LayerNorm γ 1, LayerScale γ 1e-6 — which the other threekinds already do, so this flag only has to reach the weights. It must match the JAX reference's setting (jax/MainConvNeXtImagenet.leansetscnxInit := true), or a pair run compares two inits.Separate from
vitInit, deliberately — the same splitTrainConfigkeeps on the JAX side: the two specs disagree on the conv path. timm's ViT leaves the patch embed on PyTorch'sU(±1/√fan_in)while ConvNeXt trunc-normals its convs like everything else. One boolean cannot express both, so settingvitInitfor a ConvNeXt would put the stem back at the wrong width.Variance-matched, not distribution-matched:
F32.heInitsums three uniforms (Bates-3, ≈normal) where the reference drawstrunc_normal. At σ = 0.02 that truncation is inert —trunc_normal_'s bounds are ABSOLUTE ±2, i.e. ±100σ, so nothing is ever truncated and the reference is effectively a plain normal. What remains is Bates-3's slightly lighter tail at equal σ, the same deliberate gap every other net carries.Like
vitInit, NOT display-only and needs no re-render: init is host-side, so no committed artifact moves. Off by default — every other net keeps its seed reproducibility. - cnxInit : Bool
ConvNeXt
_init_weights— the verified peer ofTrainConfig.cnxInit.trunc_normal_(std=0.02)on every conv and the head; biases 0, LayerNorm γ 1, LayerScale γ 1e-6 — which the other threekinds already do, so this flag only has to reach the weights. It must match the JAX reference's setting (jax/MainConvNeXtImagenet.leansetscnxInit := true), or a pair run compares two inits.Separate from
vitInit, deliberately — the same splitTrainConfigkeeps on the JAX side: the two specs disagree on the conv path. timm's ViT leaves the patch embed on PyTorch'sU(±1/√fan_in)while ConvNeXt trunc-normals its convs like everything else. One boolean cannot express both, so settingvitInitfor a ConvNeXt would put the stem back at the wrong width.Variance-matched, not distribution-matched:
F32.heInitsums three uniforms (Bates-3, ≈normal) where the reference drawstrunc_normal. At σ = 0.02 that truncation is inert —trunc_normal_'s bounds are ABSOLUTE ±2, i.e. ±100σ, so nothing is ever truncated and the reference is effectively a plain normal. What remains is Bates-3's slightly lighter tail at equal σ, the same deliberate gap every other net carries.Like
vitInit, NOT display-only and needs no re-render: init is host-side, so no committed artifact moves. Off by default — every other net keeps its seed reproducibility. - dwFanK2 : Bool
Depthwise fan = k², the JAX reference's rule (
jax/Jax/Codegen.lean's depthwise emitters use fank·k). The default rank-4 rule inmkParamis He fan-OUT2/(dims[0]·k²), and for a depthwise kernel(C, 1, k, k)dims[0]is C, so without this flag every depthwise conv starts √C narrower than the reference's (6–31× on MobileNetV2). Stem, 1×1s and the dense head already match.BN follows every depthwise, so the forward is unaffected; the step is not. A BN-followed weight's gradient scales as 1/‖W‖, and under TF-RMSProp at ε = 1.0 the update is ≈ g, so the relative step scales as 1/‖W‖² until weight decay equilibrates the norm.
Detection is
dims[1] = 1 ∧ dims[0] > 1, i.e. a depthwise kernel; a 1-input-channel stem would match too, and no net that sets this flag has one. Host-side, no re-render. Off in the structure; the three depthwise ImageNet drivers (MobileNetV2, MobileNetV4, EfficientNet-B0) set it since 2026-10-07, so their verified runs before that date reproduce from their seed underLEAN_MLIR_DW_FAN_K2=0, and=1turns it on for a driver that leaves it off.tests/init_parity_audit.pymeasures init parity per tensor against each net's JAX reference. - zeroGammaInit : Bool
Zero-γ on every residual-closing BN — the JAX reference's ResNet-50 init. Each bottleneck's last BN γ is init kind 4 (
VLayer.bottleneckStage); the reference starts it at 0 (emitConvBnInit … (zeroGamma := true), timm'szero_init_last=True, torchvision'szero_init_residual), so every residual branch starts as the identity. Without the flag kind 4 is γ = 1, as every verified R50 run before 2026-10-05 trained (planning/init_parity.md§2c, the audit's 16/161). Only the 16 kind-4 tensors move, each on its own seed. Host-side, no re-render. Off by default, so every other driver and every gate that inits throughmkParamis byte-identical;LEAN_MLIR_ZERO_GAMMA=0|1overrides it at launch. - seFanOutInit : Bool
Squeeze-excite FCs at the reference's fan — EfficientNet's init. The verified net carries each SE squeeze/excite as a rank-2 dense
[in,out](init kind 7,VLayer.mbConvSE); the JAX reference emits them as 1×1 convs under TF'svariance_scaling(2, fan_out), i.e.U(±√(6/out)), variance 2/out. Glorot's 2/(in+out) is a fraction of that std on every reduce FC ([480,20]: 2/500 vs 2/20), and unlike a BN-followed conv the SE FC is not scale-invariant — it feeds a sigmoid gate, so a narrow reduce FC starts every gate near σ(0) = 0.5 (planning/init_parity.md§2b, the audit's 17/213;imagenet_parity.mdD5). Only the 17 kind-7 tensors move, each on its own seed. Host-side, no re-render. Off by default, so the Imagenette B0 driver and every gate are byte-identical; the ImageNet B0 driver sets it (2026-10-07), andLEAN_MLIR_SE_FAN_OUT=0|1overrides it at launch. - headTfUniform : Bool
TF TPU EfficientNet's classifier init — the verified peer of
TrainConfig.initHeadTfUniform.efficientnet_model.py'sdense_kernel_initializerdraws the headU(±1/√shape[1])on the Keras(in, units)kernel, so the range is set by the CLASS count:U(±1/√1000)at ImageNet, variance 1/3000 against Glorot's 2/(1280 + 1000) (imagenet_parity.mdD6). Applied to the kind-0 rank-2 parameter — B0's one dense; its SE FCs are kind 7 — at variance1/(3·out)throughF32.heInit(variance-matched, as every init here). Host-side, no re-render. Off by default; the ImageNet B0 driver sets it (2026-10-12), andLEAN_MLIR_HEAD_TF_UNIFORM=0|1overrides it at launch. - truncNormalSigma : Float
TF-slim's truncated-normal init, σ on EVERY kernel — the verified peer of
TrainConfig.initTruncNormal.mobilenet.training_scopegivesslim.conv2d,slim.separable_conv2dandslim.fully_connectedonetruncated_normal_initializer(stddev)(0.09 for MobileNetV2), so the stem, the 1×1s, the depthwise kernels and the logits all start at that width, biases at 0; it replaces the fan rules (dwFanK2included) on every kind-0 and kind-7 weight. TF redraws beyond ±2σ, so the drawn std is0.8796·σ;mkParamis variance-matched to THAT (F32.heInitat0.8796·σ), as every other init here is. 0 keeps the fan rules, byte for byte;LEAN_MLIR_TRUNC_NORMAL=0|1overrides at launch. The ImageNet MobileNetV2 driver sets 0.09 (2026-10-12, planning/recipe_prep.md §6). - 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 PyTorchmomentum = 0.1is0.9here. That field's docstring carries the per-net table and the timm audit; keep the two in step, since a reference ↔ verified gap here is invisible in every loss curve.Like
vitInitand unlikelr, this is NOT display-only and needs no re-render: the graph emits raw per-layer BATCH stats and the host EMAs them (F32.emabelow), so the decay never enters the MLIR. Under gradient accumulation the driver compensates tobnMomentum^(1/k)per micro-batch, matching the reference's generated_bn. - 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 PyTorchmomentum = 0.1is0.9here. That field's docstring carries the per-net table and the timm audit; keep the two in step, since a reference ↔ verified gap here is invisible in every loss curve.Like
vitInitand unlikelr, this is NOT display-only and needs no re-render: the graph emits raw per-layer BATCH stats and the host EMAs them (F32.emabelow), so the decay never enters the MLIR. Under gradient accumulation the driver compensates tobnMomentum^(1/k)per micro-batch, matching the reference's generated_bn.
Instances For
The 95% Wilson score interval on an accuracy, as lo–hi in percentage points.
It sizes the digits an eval line can claim: on a 3,925-image validation set at p ≈ 0.85 the 95% half-width is ±1.11 pt.
Wilson, not the normal approximation p ± z·√(p(1−p)/n): the normal form collapses to ±0 at
correct = 0 or correct = n, which would print a certain 0% on a run that scored nothing.
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.
Instances For
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 for its three consumers — trainAdamSched's loop, its startup banner, and
the fp8 trainer — so the weight the banner prints is the weight the loop applies.
The == 0.99 arm returns the double 0.01 rather than 1.0 - 0.99
(= 0.010000000000000009 — different bits, 9e-16 relative), so every net that leaves
bnMomentum at its default keeps the exact weight 0.01.
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 kept 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
Instances For
Total float count across all params.
Equations
Instances For
Packed x input shape [batch, d0].
Instances For
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 records the collisions: each was between a pair of markers meeting (one marker's
spelling inside another's) rather than between a new marker and an old one.
They are defined once, here, so that the test, trainAdamSched and scoreCheckpoint all read
the same predicate: a gate on a copy of a definition is not a gate on the definition.
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.
EMA shadow — a FOURTH [θ|m|v|ema] blob region, 5 scalars not 3.
Equations
- VerifiedVariant.emaOn v = v.startsWith "ema"
Instances For
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".
Equations
- VerifiedVariant.rmsOn v = v.contains "rms"
Instances For
Stochastic depth — N extra tensor<Bxf32> scale inputs.
The marker is drop and not sd because rms ++ dp spells rmsdp, which contains
"sd".
Equations
- VerifiedVariant.sdOn v = v.contains "drop"
Instances For
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.
Equations
- VerifiedVariant.cdOn v = v.contains "do"
Instances For
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 = v.contains "acc"
Instances For
LAMB — the per-tensor trust ratio (OptRecipe.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 = v.contains "lamb"
Instances For
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 = v.contains "bce"
Instances For
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.
Instances For
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.
The fifth region lets accumulation and EMA run together (RSB-A2/A1's recipe sets both
gradAccumSteps := 4 and useEMA := true).
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 a single-axis checkpoint reads at the same index either way.
The reverse order would re-home every ema* blob.
Equations
Instances For
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
Instances For
The blob index of the EMA shadow region, or none when the variant has no shadow.
It is not the literal 3: 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
Instances For
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
Instances For
The eval forward's suffix for a variant rendered at a non-default BatchNorm ε: _eps0001 for
an …eps0001… variant (ε = 1e-3), empty otherwise. ε is baked into every BN site, so a train
step at one ε must be scored through an eval graph at the same ε —
<slug>_fwd_eval_eps0001.mlir, entry @<slug>_fwd_eval_eps0001 — while every checkpoint
trained at 1e-5 keeps scoring through <slug>_fwd_eval.mlir.
Equations
- VerifiedVariant.evalTag v = match v.splitOn "eps" with | head :: rest :: tail => "_eps" ++ String.ofList (List.takeWhile Char.isDigit rest.toList) | x => ""
Instances For
The exact GELU x · Φ(x) rather than the tanh approximation — the renderers' geluMarker,
erf after drop or the ε marker and before bf16 (emadpwxclipdroperfbf16,
emadp128x4wxclipdropeps0000001erfbf16). Unlike the layout markers it changes no arity, type
or region. The driver reads it for one thing: which forward the variant scores through
(fwdStem), with checkGeluWorld refusing a train step and a forward at different forms.
Equations
- VerifiedVariant.erfOn v = v.contains "erf"
Instances For
The forward a variant scores through when it has none of its own (<slug>_<variant>_fwd):
<slug>_fwd, or <slug>_erf_fwd for an exact-GELU variant. The GELU is baked into every
block of the forward, so the tanh <slug>_fwd is a different net from the one an …erf…
train step trains.
Equations
Instances For
iree-compile one .mlir → .vmfb, surfacing failures. Skips when the .vmfb is already
newer than the .mlir (a content-stable cache): avoids the ~minutes-long 224² recompile, and
lets two same-net runs share one GPU-pair safely — they only read the cached vmfb (concurrent
reads are fine; it's the concurrent writes of an identical compile that would race).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Open a session for one Lean-emitted graph, on whichever backend this binary dlopened.
- XLA — hand the
.mlirstraight to PJRT, which compiles it in-process. Nothing is written to disk. - IREE —
iree-compilethe.mlirto a cache file first, then load that.
The cache path is derived from mlirPath rather than passed in, so two
graphs cannot land on one cache file — the failure the target scoping below
prevents from the other direction.
Both backends consume the same verified_mlir/*.mlir — the emitter, the
spec, and the train-step ties are identical. Only the trusted lowerer differs.
Equations
- One or more equations did not get rendered due to their size.
Instances For
mkSession for the eval forward at replicas devices, the sharded-inference session
(LowererSession.createDp, outputs gathered). replicas ≤ 1 IS mkSession, so a single-GPU
run and every IREE run are unchanged. Past 1 it is XLA-only and says so rather than quietly
evaluating on one card.
Equations
- One or more equations did not get rendered due to their size.
Instances For
One plain SGD step of net's rendered @<slug>_train_step on the params-only theta,
returning the params-only update — for the drivers that keep no loss (PGD, smoothing). A
lossSlot render takes a trailing scalar and returns it written; the slot is appended here
and dropped from the result, so the caller never sees it.
Equations
- One or more equations did not get rendered due to their size.
Instances For
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, variance2/(oc·kH·kW) - rank-2 dense matrix
[in, out]→ Glorot, variance2/(in + out) - γ = 1 (kind 1), β / bias = 0 (kind 2), layer scale = 1e-6 (kind 3)
- zero-γ (kind 4: a bottleneck's residual-closing BN γ) → 0 under
zeroGamma, 1 otherwise (planning/init_parity.md§2c) - embedding (kind 5: ViT's CLS token and positional embedding) → σ = 0.02 under
vitInit, 0 otherwise (planning/init_parity.md§2d, §3a) - SE FC (kind 7: EfficientNet's squeeze / excite dense
[in,out]) → variance2/outunderseFanOut, the reference's 1×1-conv fan-out; Glorot otherwise (planning/init_parity.md§2b) truncNormal > 0(TF-slim'struncated_normal_initializer(stddev),VerifiedConfig.truncNormalSigma) → every weight at std0.8796·σ, the std of a normal truncated at ±2σ, before any fan rule; biases stay 0
The weight variances are those of
jax/Jax/Codegen.lean:
uniform(±√(6/fan_out)) for convs (torchvision's
kaiming_normal_(mode='fan_out', nonlinearity='relu') convention, emitConvBnInit) and
uniform(±√(6/(fan_in+fan_out))) for dense (emitDenseInit); kind 3 is the reference's
ConvNeXt layer-scale init, 1e-6. Kind 3 must be matched explicitly: the _ branch is the weight rule,
which is the wrong answer for a per-channel scale. cnxInit and vitInit override the weight
variance (see VerifiedConfig); heFanIn switches to fan-in for one gate.
Init is host-side: no verified_mlir/ file mentions it, so changing it moves no committed
artifact.
Variance-matched, not distribution-matched: F32.heInit sums three uniforms (Bates-3,
≈ normal) where JAX draws one uniform. torchvision itself uses a normal here, so neither side
is canonical on that axis.
Equations
- One or more equations did not get rendered due to their size.
- mkParam seed dims 1 vitInit biasSigma heFanIn cnxInit dwFanK2 zeroGamma seFanOut truncNormal headTfUniform = F32.const (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize 1.0
- mkParam seed dims 3 vitInit biasSigma heFanIn cnxInit dwFanK2 zeroGamma seFanOut truncNormal headTfUniform = F32.const (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize 1e-6
- mkParam seed dims 2 vitInit none heFanIn cnxInit dwFanK2 zeroGamma seFanOut truncNormal headTfUniform = F32.const (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize 0.0
- mkParam seed dims 2 vitInit (some s) heFanIn cnxInit dwFanK2 zeroGamma seFanOut truncNormal headTfUniform = F32.heInit seed.toUSize (Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 dims).toUSize s
Instances For
The ImageNet batch shim #
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.
Read up to len bytes from h into buf, appending at buf.size and never touching its
capacity — which is the reason it exists. IO.FS.Handle.read allocates its result on the
CALLING thread, so a batch read on a pool thread and dropped on the main thread is freed by a
thread that does not own it. The runtime's allocator (mimalloc) answers a cross-thread free of
a huge block with madvise(MADV_FREE), not a release: the pages stay in RSS as LazyFree
until the kernel is under pressure; at ImageNet batch sizes that fills host memory within an
epoch and slows every step through reclaim. With this primitive the main thread allocates AND frees every batch buffer; the pool thread
only fills it, and the block is recycled in place step after step.
buf must be the only reference (rc 1, or −1 once it has crossed a Task boundary). The C
side refuses a shared buffer rather than copying it, because a silent copy would put the
allocation straight back on the reading thread. Returns the buffer with its size advanced by
the bytes read; 0 bytes means end of stream, as for Handle.read.
Fill buf with EXACTLY n more 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.
Instances For
Read EXACTLY n bytes into a fresh buffer of exactly that capacity, allocated HERE — on the
calling thread, which is what readInto is about. Allocating the exact capacity up front
avoids ByteArray.append's doubling (a 308 MB batch would cost a 616 MB buffer and a copy).
Equations
- readExact h n = readExactInto h (ByteArray.emptyWithCapacity n) n
Instances For
The stdio shape every shim child is spawned with. It has a name because IO.Process.Child is
INDEXED by its config: without a concrete one the child cannot be stored in a structure, and
without storing it the trainer holds a pipe it can neither kill nor reap.
Equations
- ShimCfg = { stdin := IO.Process.Stdio.null, stdout := IO.Process.Stdio.piped }
Instances For
A live shim producer: the child process AND its stdout pipe.
Holding the child lets the trainer reap it (no <defunct> process for the life of the run)
and replace a degraded producer mid-run; under the round-robin read one slow loader paces
the whole run.
- child : IO.Process.Child ShimCfg
- child : IO.Process.Child ShimCfg
- h : IO.FS.Handle
- h : IO.FS.Handle
some mswhen the producer came up on the uint8 wire (v5/v6): the six float32 normalize constants from its preamble, whichreadShimBatchhands tou8Norm.
Instances For
The uint8 wire's consumer half (lean_mlir_u8_norm): buf holds B uint8 HWC images at its
front and the capacity of the float batch; it comes back as B float32 CHW images,
((float)u − mean[c]) / std[c] — the shim's own normalize and transpose, op for op — in place,
so the main thread still allocates and frees every batch buffer (see readInto).
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.
Instances For
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.
Instances For
One shim batch, tolerating a SHORT FINAL BATCH — the validation-split reader.
The shim sets drop_remainder=training, so the validation split ends in a partial batch
(ImageNet's 50,000 images at 256 are 195 full batches and an 80-image tail) that readExact
refuses by construction. This reader accepts it, so every image is scored — timm's
validate.py denominator, not 49,920.
Returns (img, lbl, rows) where rows ≤ batch, and rows = 0 means the stream ended cleanly.
Each batch starts with a 4-byte row count (read, not inferred: a pipe does not preserve
write boundaries), then exactly rows labels and rows images, in wire order. A short read
of either block is a torn write and throws instead of returning a silently short batch.
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.
Instances For
Where evalScore reads its rows from.
- held
(img lbl : ByteArray)
: EvalRows
The whole split held in RAM — MNIST, CIFAR, Imagenette.
- stream
(hs : Array ShimProc)
(shimBatch flat offset : Nat)
(dropTail : Bool)
: EvalRows
ImageNet: streamed per pass from
hs.sizebatch-block producers (spawnValStream),shimBatchrows per shim batch,flatfloats per image, read round-robin — global batchkfrom producer(k + offset) % n, whereoffsetis 0 except under the gate's order fault.dropTaildrops a partial final batch (LEAN_MLIR_EVAL_BATCHSTATS, or the gate's tail fault) instead of refusing the short pass.
Instances For
The streamed reader's carry. Shim batches arrive at shimBatch rows and an eval invoke wants
gB = R × evalBs of them — one shim batch for ConvNeXt at 4 × 64, four for ViT at 4 × 256, a
quarter of one at R = 1, bs 64 — so rows are carried across pulls.
- img : ByteArray
- img : ByteArray
- lbl : ByteArray
- lbl : ByteArray
- rows : Nat
- rows : Nat
- next : Nat
- next : Nat
- ended : Bool
- ended : Bool
- total : Nat
- total : Nat
Instances For
Spawn the per-pass val producers: n batch-block producers of
the validation split at shimBatch rows — fresh for every pass and reaped by reapValStream, so
nothing accumulates across epochs and the closed pipe stays the end-of-pass marker. flat is the
EVAL width (RSB-A3 trains at 160² and evaluates at 224²), read off the artifact by the caller.
The gate-only fault knobs, in the PJRT_FFI_FAULT style — controls that must go red:
LEAN_MLIR_VAL_FAULT=order starts the round-robin on producer 1 (every image scored against
another image's label: same count, different bitmap), =tail drops the 80-row tail (a
49,920-image pass). scripts/gates/streamed_val_gate.sh.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The eval pass: score nEval images through a forward session, replicas × evalBs per
invoke. (top-1 correct, top-5 correct, images scored, per-image top-1 bitmap). The bitmap is
filled only when wantBits, in eval order, one byte per image. rows is where the images come
from: the held split, or ImageNet's per-pass stream (EvalRows).
One copy, shared by the per-epoch eval in trainAdamSched and by scoreCheckpoint, because
scripts/gates/sharded_eval_gate.sh compares 1 replica against N through score-checkpoint,
and a gate on a copy of the loop says nothing about the loop that runs.
The ragged tail: 50,000 is not a multiple of 4 × 64 = 256:
it is 195 full invokes and an 80-image tail, which fills replica 0 and a quarter of replica 1,
while replicas 2 and 3 score pure padding. F32.sliceImagesPad zero-pads the invoke to the
global batch, the shim gathers the logits back in ROW ORDER, and only
min gB (nEval − bi·gB) rows are scored. The pad rows are computed and never read.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Load the train + eval splits for a dataset. Returns
(trainImg, trainLbl, nTrain, evalImg, evalLbl, nEval, trainPix, crop?) where
trainPix is the stored per-example width of the training images (256² for
Imagenette, d0 otherwise) and crop? requests the 256²→224² center-crop.
evalOnly skips the TRAIN split entirely — for scoreCheckpoint, which never touches it.
It is inert on .imagenet, whose train split is never preloaded (it streams off the shim),
and worth having on the others: Imagenette's is 9,469 × 256² × 3 f32 = 7.4 GB read and held for a job that only scores 3,925 val images. nTrain comes back 0
under it, so a caller that starts using it gets a division rather than a plausible epoch.
Equations
- One or more equations did not get rendered due to their size.
Instances For
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. $LEAN_MLIR_CKPT_TAG appends a run-scoped suffix — without it every pass of the same
(net, variant, backend) shares one path, so parallel passes clobber each other's blob and a
later pass resumes from an earlier one's finished run.
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.
Instances For
Print the startup banner with %LOWERER% resolved to the lowerer that actually ran.
The blurb carries a placeholder rather than a transport, so the banner states the run, not the build; every print site goes through here.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Train a VerifiedNet end-to-end on its proof-rendered StableHLO: compile both
MLIRs → lowerer sessions → load data → spec init (mkParam) → 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.
Instances For
The LayerNorm-ε world check — the LN peer of the BN-world check in trainAdamSched.
A no-BN net evaluates through <slug>_<variant>_fwd.mlir when it exists and otherwise
falls back to <slug>_fwd.mlir. A train step rendered at a non-default LN ε (ViT's
eps0000001 variants, DeiT's 1e-6) whose per-variant forward is missing would TRAIN at 1e-6
and SCORE at 1e-5 — a different net, and a plausible number. The ε prints as a dense<…>
literal, so compare the set of LN-ε spellings the two artifacts carry. 1.0e-6 is
unambiguous: across verified_mlir/ only LayerNorm emits it (the clip's 1e-6 prints
0.000001), so ConvNeXt (1e-6 on both sides) and every 1e-5 net pass unchanged.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The GELU-form world check — checkLnEpsWorld's peer for the form of the GELU.
The spec language does not carry the form (it is a render parameter, GeluForm), so the
variant name is the only place the driver can read it and the artifacts are the only place it
is true. An exact-GELU train step scored through a tanh forward, or the reverse, is a
different net and a plausible number. Three readings must agree: the name's erf marker
(VerifiedVariant.erfOn), the train step's ops and the forward's. Across verified_mlir/
only the exact GELU emits chlo.erfc, so a net with no GELU carries none on either side and
passes unchanged.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Score a checkpoint, standalone — the eval half of trainAdamSched with no training in
front of it. It scores any saved checkpoint, not only the weights live during training — the
verified peer of the JAX side's eval_*_full50k.py. No new MLIR: every piece is the eval
half's own:
| need | reused |
|---|---|
| read the val split | spawnValStream (ImageNet, streamed per pass) / loadData (held splits) |
| 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 |
| the eval pass | evalScore (short-tail padding, forward, scoring) |
| 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.
BN nets score through @<slug>_fwd_eval with the running statistics read from the
<ckpt>.bn companion ([running | ema_bn], written at every epoch end): the
EMA shadow with ema_bn, the live weights with the running stats — the training eval's
pairing. A checkpoint older than the companion is refused, since scoring the .bin alone would
normalise by zeros and still print a plausible percentage.
timm's test protocol: LEAN_MLIR_EVAL_SIZE / LEAN_MLIR_EVAL_CROP score through the eval
graph rendered at that size (…_fwd_eval_s<S>.mlir) with the val stream resized and cropped to
match (SHIM_EVAL_SIZE / SHIM_EVAL_CROP). scripts/parity/score_timm.sh reads the per-net values
from jax/timm_eval_protocols.json.
region is what one checkpoint cannot otherwise yield: the driver picks live-or-shadow at
TRAIN time (LEAN_MLIR_EMA_BN), 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.
Instances For
timm's CosineLRScheduler with t_in_epochs (DeiT, the RSB tiers), the rate at 0-based epoch
ep, held for the whole epoch: warmupLRInit + t·(baseLR − warmupLRInit)/warmupEpochs for
t < warmupEpochs, else minLR + ½(baseLR − minLR)(1 + cos(π t / epochs)) — the cosine over ALL
epochs counted from epoch 0 (warmup_prefix = False). The reference emits the same two lines
(TrainConfig.schedPerEpoch);
scripts/gates/sched_timm_gate.py
checks both against timm itself, this one through
tests/TestSchedPerEpoch.lean.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Scheduled AdamW driver — threads the first/second moment buffers as one packed
[θ|m|v] param blob through the generic FFI (n_params = 3k), 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 the committed
<mlirDir>/<slug>_<variant>_train_step.mlir, rendered by LeanMlir/Proofs/Codegen/; its
optimizer ops are tied per op (for example Proofs.StableHLO.adamW_triple_faithful,
Proofs.StableHLO.lamb_triple_faithful, Proofs.StableHLO.mom_pair_faithful).
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). It is a
driver item, not a render one, because lr is already a runtime operand. Default 0.0
keeps cosine.
RMSProp variants also need a different initial state (the mean-square slot starts at
1.0, VerifiedVariant.rmsOn), handled in the body.
Equations
- One or more equations did not get rendered due to their size.
Instances For
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, mkParam-init train above. Only the linear classifier uses this;
shares mkSession / loadData with the main driver.
Equations
- One or more equations did not get rendered due to their size.
Instances For
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
render-tie model of Proofs/Float/E4M3Fold.lean). Eval runs the fp32
master through @<slug>_fwd (the "fp32-infer" accuracy of the fp8-trained
model, mirroring scripts/demos/mnist_e4m3_demo.py).
Run: lake exe mnist-linear-e4m3-verified data
Equations
- One or more equations did not get rendered due to their size.
Instances For
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: 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 would need in-graph E4M3 ops, not host byte-prep. So this is fp8 weights + fp8 input, fp32 intermediates. Eval runs the fp32 master.
Run: lake exe mnist-mlp-e4m3-verified data
Equations
- One or more equations did not get rendered due to their size.
Instances For
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.
Instances For
Lower to the runtime VerifiedNet the driver consumes.
Equations
- One or more equations did not get rendered due to their size.