Unified training loop for any spec that uses the Adam codegen path.

Each Main*Train.lean used to inline ~250 lines of identical code: generate MLIR → compile vmfbs → init params → for-each-epoch { shuffle, for-each-batch { trainStepAdamF32, EMA, log }, val eval } → save. The body was the same for ResNet/MobileNet/EfficientNet/ViT/VGG modulo a handful of name strings.

NetSpec.train is the extracted function. A trainer is now:

def main (args : List String) : IO Unit :=
  resnet34.train resnet34Config (args.head?.getD "data/imagenette")

The 250 lines collapse to one call. Adding a new architecture means defining the spec + a TrainConfig value — no copy-paste plumbing.

def NetSpec.buildPrefix (spec : NetSpec) : String

File-path prefix for the generated MLIR / vmfb / saved-params files associated with this spec. Uses the sanitized spec name so adding a new spec automatically gets a unique non-colliding prefix without the trainer author having to pick one.

Equations

• spec.buildPrefix = ".lake/build/" ++ spec.sanitizedName ++ if spec.buildTag.isEmpty = true then "" else "_" ++ MlirCodegen.sanitize spec.buildTag

def NetSpec.withBuildTag (spec : NetSpec) (tag : String) : NetSpec

Set buildTag, so an ablation arm owns its own artifacts. Reads as (net.withBuildTag "ce").train … at the call site. (Not buildTag — that name belongs to the field itself.)

Equations

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

def NetSpec.trainFnName (spec : NetSpec) : String

The fully qualified <module>.<func> name the train step's main function lives at — what trainStepAdamF32 wants as its fnName. Mirrors how MlirCodegen.generateTrainStep builds the module name from the third argument it's given (jit_<sanitized>_train_step).

Equations

• spec.trainFnName = "jit_" ++ spec.sanitizedName ++ "_train_step.main"

def NetSpec.graphArtifact (pfx suffix : String) : IO String

Path of the runnable artifact for one emitted graph, on whichever backend is ACTIVE AT RUN TIME (planning/archive/detector_pjrt_port.md).

⚠ This said "whichever backend this binary was linked against" until 2026-08-25, and as of that date nothing is linked against a backend at all: ireeLink/xlaLink are retired and every executable is on lowererLink, with ffi/lowerer.c dlopening the shim $LEAN_MLIR_LOWERER names. The dispatch below was always the real mechanism — it reads LowererSession.backendName, not the link line — so the behaviour is unchanged; only the sentence was wrong.

• IREE — iree-compile turns {pfx}_{suffix}.mlir into a .vmfb, and that .vmfb is what LowererSession.create loads.
• XLA — PJRT compiles the .mlir in-process, so no .vmfb is ever produced and the .mlir is the artifact.

Every pathExists guard and LowererSession.create call below routes through this. The guards are the reason it exists: they are what decide whether eval runs at all, and with a hardcoded .vmfb path they go silently false on XLA — training would look perfectly healthy and simply stop reporting val metrics, which is a much worse failure than not loading.

Equations

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

def NetSpec.compileVmfbs (spec : NetSpec) (cfg : TrainConfig) (useSeg : Bool := false) : IO String

Generate forward / eval-forward / train-step MLIR for spec, write them to .lake/build/<sanitized>_*.mlir, and compile each to a .vmfb. Cached on the MLIR content + IREE backend so a second run with no codegen changes skips iree-compile entirely (saves ~10-15 min for ResNet-sized models). Returns the path to the train-step vmfb.

useSeg = true switches the train-step codegen to per-pixel softmax CE (segmentation), where labels are an int32 [B, H, W] tensor instead of a [B] class index. Mutually exclusive with soft-label / focal / label-smoothing flows.

Equations

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

@[implicit_reducible]

instance NetSpec.instInhabitedDatasetIO : Inhabited NetSpec.DatasetIO✝

Equations

• NetSpec.instInhabitedDatasetIO = { default := NetSpec.instInhabitedDatasetIO.default }

def NetSpec.runTraining (spec : NetSpec) (cfg : TrainConfig) (ds : DatasetKind) (dataDir : String) (sess : LowererSession) : IO Unit

Adam + cosine-LR + running-BN-stats training loop, generic over DatasetKind. The dataset specifies how to load the train/val data and what augmentation to apply per batch; everything else (init, optimizer, BN EMA, val eval, save) is identical across datasets.

spec must have been compiled via compileVmfbs first.

def NetSpec.evalOnly (spec : NetSpec) (cfg : TrainConfig) (ds : DatasetKind) (dataDir : String) : IO Unit

Eval-only mode: skip training entirely, load saved params + bn_stats, run the eval forward on val data, print accuracy. Used to re-eval existing checkpoints against a fixed eval pipeline (e.g. after the Imagenette centerCrop bug was fixed). Requires compileVmfbs to have produced (or cached) the _fwd_eval.vmfb already.

Equations

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

def NetSpec.train (spec : NetSpec) (cfg : TrainConfig) (dataDir : String) (ds : DatasetKind := DatasetKind.imagenette) : IO Unit

End-to-end: compile all three vmfbs, load the train-step session, and run the training loop on the chosen dataset. The high-level entry point that every Main*Train.lean now calls.

LEAN_MLIR_EVAL_ONLY=1 short-circuits to evalOnly — skips training, loads the saved checkpoint, runs eval. Used for re-eval after a fixed eval pipeline.

Defaults to Imagenette so existing trainers don't have to change.