Documentation

LeanMlir.Verified.Train

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:

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 #

structure VerifiedNet :

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 contents scripts/regen_verified_mlir.sh check pins to exactly the set with a literal IO.FS.writeFile "verified_mlir/…" writer in LeanMlir/Proofs/Codegen/.

    The width/batch sweep nets (mlpG, cnnG, cifar8BnG) set .lake/build instead: 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 of verified_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 contents scripts/regen_verified_mlir.sh check pins to exactly the set with a literal IO.FS.writeFile "verified_mlir/…" writer in LeanMlir/Proofs/Codegen/.

    The width/batch sweep nets (mlpG, cnnG, cifar8BnG) set .lake/build instead: 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 of verified_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 the m.<slug>_{train_step,fwd} funcs.

  • 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 = random weight (conv: He fan-out; dense: Glorot; see mkParam), 1 = ones (γ), 2 = zeros (β / bias), 3 = 1e-6 (layer scale).

  • specs : Array (Array Nat × Nat)

    (dims, initKind) per param, in func-arg order — the matching XLayout.specs. initKind: 0 = random weight (conv: He fan-out; dense: Glorot; see mkParam), 1 = ones (γ), 2 = zeros (β / bias), 3 = 1e-6 (layer scale), 4 = zero-γ (0 under zeroGammaInit, else 1), 5 = embedding (σ 0.02 under vitInit, else 0), 7 = SE FC (var 2/out under seFanOutInit, 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.

  • Dataset / loader selector.

  • 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.

  • 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: true for the mlp and cnn renders, 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.mlir return the trailing report-only %loss scalar?

    A per-render fact, not a driver-wide one: true for the mlp and cnn renders, 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.

  • 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.

  • 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. 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_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, 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.

  • dropKeeps : Array Float

    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_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, like toSpecs == XLayout.specs. tests/TestDropPathRamp.lean is the #guard that pins the two; Verified.Spec sits downstream of this file, so the renderer cannot share the definition by import without inverting the dependency.

  • dropoutKeep : Option (Float × Nat)

    Classifier dropout — (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).

    The width is here because the mask is per-element, which is the 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.

  • dropoutKeep : Option (Float × Nat)

    Classifier dropout — (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).

    The width is here because the mask is per-element, which is the 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.

    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 .imagenet net refuses at spawn, because a wrong shim fails silently — 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.

  • 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.

    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 .imagenet net refuses at spawn, because a wrong shim fails silently — 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.

Instances For
    structure VerifiedConfig :

    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, like LEAN_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, like LEAN_MLIR_MAX_EPOCHS.

    • vitInit : Bool

      timm/DeiT ViT weight init — the verified peer of TrainConfig.vitInit, i.e. of the deit-init recipe of the JAX reference. 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.

    • vitInit : Bool

      timm/DeiT ViT weight init — the verified peer of TrainConfig.vitInit, i.e. of the deit-init recipe of the JAX reference. Every weight at σ = 0.02 except the patch-embed conv on PyTorch's U(±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 of TrainConfig.cnxInit. trunc_normal_(std=0.02) on every conv and the head; biases 0, LayerNorm γ 1, LayerScale γ 1e-6 — which the other three kinds already do, so this flag only has to reach the weights. It must match the JAX reference's setting (jax/MainConvNeXtImagenet.lean sets cnxInit := true), or a pair run compares two inits.

      Separate from vitInit, deliberately — the same split TrainConfig keeps on the JAX side: the two specs disagree on the conv path. timm's ViT leaves the patch embed on PyTorch's U(±1/√fan_in) while ConvNeXt trunc-normals its convs like everything else. One boolean cannot express both, so setting vitInit for a ConvNeXt would put the stem back at the wrong width.

      Variance-matched, not distribution-matched: F32.heInit sums three uniforms (Bates-3, ≈normal) where the reference draws trunc_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 of TrainConfig.cnxInit. trunc_normal_(std=0.02) on every conv and the head; biases 0, LayerNorm γ 1, LayerScale γ 1e-6 — which the other three kinds already do, so this flag only has to reach the weights. It must match the JAX reference's setting (jax/MainConvNeXtImagenet.lean sets cnxInit := true), or a pair run compares two inits.

      Separate from vitInit, deliberately — the same split TrainConfig keeps on the JAX side: the two specs disagree on the conv path. timm's ViT leaves the patch embed on PyTorch's U(±1/√fan_in) while ConvNeXt trunc-normals its convs like everything else. One boolean cannot express both, so setting vitInit for a ConvNeXt would put the stem back at the wrong width.

      Variance-matched, not distribution-matched: F32.heInit sums three uniforms (Bates-3, ≈normal) where the reference draws trunc_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 fan k·k). The default rank-4 rule in mkParam is He fan-OUT 2/(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 under LEAN_MLIR_DW_FAN_K2=0, and =1 turns it on for a driver that leaves it off. tests/init_parity_audit.py measures 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's zero_init_last=True, torchvision's zero_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 through mkParam is byte-identical; LEAN_MLIR_ZERO_GAMMA=0|1 overrides 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's variance_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.md D5). 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), and LEAN_MLIR_SE_FAN_OUT=0|1 overrides it at launch.

    • headTfUniform : Bool

      TF TPU EfficientNet's classifier init — the verified peer of TrainConfig.initHeadTfUniform. efficientnet_model.py's dense_kernel_initializer draws the head U(±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.md D6). Applied to the kind-0 rank-2 parameter — B0's one dense; its SE FCs are kind 7 — at variance 1/(3·out) through F32.heInit (variance-matched, as every init here). Host-side, no re-render. Off by default; the ImageNet B0 driver sets it (2026-10-12), and LEAN_MLIR_HEAD_TF_UNIFORM=0|1 overrides it at launch.

    • truncNormalSigma : Float

      TF-slim's truncated-normal init, σ on EVERY kernel — the verified peer of TrainConfig.initTruncNormal. mobilenet.training_scope gives slim.conv2d, slim.separable_conv2d and slim.fully_connected one truncated_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 (dwFanK2 included) on every kind-0 and kind-7 weight. TF redraws beyond ±2σ, so the drawn std is 0.8796·σ; mkParam is variance-matched to THAT (F32.heInit at 0.8796·σ), as every other init here is. 0 keeps the fan rules, byte for byte; LEAN_MLIR_TRUNC_NORMAL=0|1 overrides 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 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 reference ↔ verified 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.

    • 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 reference ↔ verified 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.

    Instances For
      def wilson95 (correct nEval : Nat) :

      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
        Instances For

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

          Equations
          Instances For

            Packed shape descriptors for the FFI (see packShapes).

            Equations
            Instances For

              Total float count across all params.

              Equations
              Instances For

                Packed x input shape [batch, d0].

                Equations
                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
                  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
                    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
                      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
                        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
                          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
                            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
                              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
                                        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
                                          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
                                            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
                                                def compileVmfb (mlirPath outPath : String) :

                                                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
                                                  def mkSession (mlirPath : String) :

                                                  Open a session for one Lean-emitted graph, on whichever backend this binary dlopened.

                                                  • 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, 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
                                                    def mkSessionDp (mlirPath : String) (replicas : Nat) :

                                                    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
                                                      def VerifiedNet.sgdStep (net : VerifiedNet) (sess : LowererSession) (fn : String) (xb theta yb : ByteArray) (bs : Nat) :

                                                      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
                                                        def mkLabels (bs off nc : Nat) :

                                                        bs labels, class (i + off) % nc, as the driver's int32 records — the shard labels the sync-BN, shard and tie gates feed.

                                                        Equations
                                                        • One or more equations did not get rendered due to their size.
                                                        Instances For
                                                          def mkParam (seed : Nat) (dims : Array Nat) (kind : Nat) (vitInit : Bool := false) (biasSigma : Option Float := none) (heFanIn cnxInit dwFanK2 zeroGamma seFanOut : Bool := false) (truncNormal : Float := 0.0) (headTfUniform : Bool := false) :

                                                          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), 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]) → variance 2/out under seFanOut, the reference's 1×1-conv fan-out; Glorot otherwise (planning/init_parity.md §2b)
                                                          • truncNormal > 0 (TF-slim's truncated_normal_initializer(stddev), VerifiedConfig.truncNormalSigma) → every weight at std 0.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
                                                          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.

                                                            @[extern lean_mlir_read_into]
                                                            opaque readInto (h : IO.FS.Handle) (buf : ByteArray) (len : USize) :

                                                            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
                                                              Instances For
                                                                @[reducible, inline]

                                                                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
                                                                Instances For
                                                                  structure ShimProc :

                                                                  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.

                                                                  Instances For
                                                                    @[extern lean_mlir_u8_norm]
                                                                    opaque u8Norm (buf : ByteArray) (B hw : USize) (ms : ByteArray) :

                                                                    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).

                                                                    def spawnShim (shimScript split : String) (batch flat seed : Nat) (shard : Option (Nat × Nat) := none) (nclasses : Nat := 0) (u8 : Bool := false) (extraEnv : Array (String × Option String) := #[]) :

                                                                    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
                                                                        def readShimBatchPartial (h : IO.FS.Handle) (batch flat : Nat) (nclasses : Nat := 0) :

                                                                        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
                                                                          inductive EvalRows :

                                                                          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.size batch-block producers (spawnValStream), shimBatch rows per shim batch, flat floats per image, read round-robin — global batch k from producer (k + offset) % n, where offset is 0 except under the gate's order fault. dropTail drops a partial final batch (LEAN_MLIR_EVAL_BATCHSTATS, or the gate's tail fault) instead of refusing the short pass.

                                                                          Instances For
                                                                            structure ValCarry :

                                                                            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.

                                                                            Instances For
                                                                              def spawnValStream (net : VerifiedNet) (flat : Nat) (n : Nat := 2) (shimBatch : Nat := 256) (extraEnv : Array (String × Option String) := #[]) :

                                                                              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
                                                                                def evalScore (sess : LowererSession) (fn : String) (params shapes : ByteArray) (rows : EvalRows) (nEval evalBs evalD0 nc replicas : Nat) (nResident gen : USize) (wantBits : Bool) :

                                                                                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
                                                                                  def loadData (net : VerifiedNet) (dataDir : String) (evalD0 : Nat := 0) (evalOnly allowU8 : Bool := false) :

                                                                                  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
                                                                                        def VerifiedNet.train (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) :

                                                                                        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
                                                                                          def checkLnEpsWorld (tsPath fwdPath : String) :

                                                                                          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
                                                                                            def checkGeluWorld (variant tsPath fwdPath : String) :

                                                                                            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
                                                                                              def VerifiedNet.scoreCheckpoint (net : VerifiedNet) (dataDir variant ckptPath : String) (region : String := "auto") :

                                                                                              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:

                                                                                              needreused
                                                                                              read the val splitspawnValStream (ImageNet, streamed per pass) / loadData (held splits)
                                                                                              the eval graphmkSession on <slug>_fwd_eval, or the _fwd chain for the LN nets
                                                                                              eval batch AND widthfwdRenderedShape, one parse of one declaration
                                                                                              the eval passevalScore (short-tail padding, forward, scoring)
                                                                                              forwardLowererSession.forwardF32
                                                                                              metricsF32.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
                                                                                                def schedPerEpochLR (baseLR minLR warmupLRInit : Float) (warmupEpochs epochs ep : Nat) :

                                                                                                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
                                                                                                  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) (expStaircase : Bool := false) (minLR : Float := 0.0) (emaWarmup : Bool := true) (schedPerEpoch : Bool := false) (warmupLRInit : Float := 0.0) (cosineFromStep0 : Bool := false) :

                                                                                                  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
                                                                                                          def VerifiedNet.trainAdamSchedE4M3 (net : VerifiedNet) (cfg : VerifiedConfig) (dataDir : String) (baseLR β1 β2 : Float) (warmupEpochs : Nat) (variant : String := "adam") (expDecayRate : Float := 0.0) :

                                                                                                          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.
                                                                                                            Instances For