Documentation

(inter, |gt|, |pred|) of the class-union region cls, read off an nc × nc confusion matrix conf[gt * nc + pred]: inter = Σ_{g∈R} Σ_{p∈R} C[g][p], |gt| = Σ_{g∈R} Σ_p C[g][p], |pred| = Σ_{p∈R} Σ_g C[g][p]. Counts stay exact Nat.

2·inter / (|gt| + |pred|), with the BraTS convention for a region the ground truth lacks: an empty prediction scores 1 and a non-empty one 0.

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.

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

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

Cache key for an MLIR string under the current IREE backend. Combines the MLIR content with the backend env var so switching IREE_BACKEND=rocm ↔ cuda invalidates the cache.

Path of the runnable artifact for one emitted graph, on whichever backend is active at run time (LowererSession.backendName). No executable is linked against a backend: every one links lowererLink, and ffi/lowerer.c dlopens the shim $LEAN_MLIR_LOWERER names.

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

The eval forward of spec over the n images of a flat f32 set (imgFloats floats each) at batch batch, the tail batch padded with the last real image: the [n, nOut] f32 logits of the real rows. The detector demos' infer dumps.

Compile one emitted graph to the artifact graphArtifact names: iree-compile on IREE; on XLA nothing, since PJRT compiles the .mlir itself in-process. false on a failed compile (the compiler's error is printed). Uncached; compileVmfbs keys its own compiles on the MLIR text.

Compile mlirPath to outPath via iree-compile, but skip the work if outPath exists and outPath ++ ".hash" matches the cache key for the current MLIR + IREE backend. Returns (compiledNewly, success).

Generate forward / eval-forward / train-step MLIR for spec and write them to .lake/build/<sanitized>_*.mlir. On IREE, also compile each to a .vmfb, cached on the MLIR content + IREE backend so a second run with no codegen changes skips iree-compile; on XLA there is no compile step. Returns the train-step artifact graphArtifact resolves (.mlir on XLA, .vmfb on IREE).

The loss is cfg.lossKindFor ds; a per-pixel kind switches the train-step codegen to the segmentation path, where labels are an int32 [B, H, W] tensor instead of a [B] class index.

Imagenette: 256×256 train, 224×224 val, random-crop-to-224 + hflip. When augment=false, falls back to deterministic center-crop 256→224 so the model still receives correctly-shaped input.

MNIST: 28×28 IDX format, no augmentation.

CIFAR-10: 5 train batches concatenated, 1 test batch, random hflip.

BraTS brain-tumour segmentation (MSD Task01_BrainTumour), 2D axial slices. 240×240 images with 4 MRI modalities as channels (FLAIR / T1w / T1gd / T2w) and 240×240 per-pixel masks (4-class: 0=background, 1=edema, 2=non-enhancing tumour, 3=enhancing tumour). The "labels" buffer is the mask buffer (240*240 bytes per record), which is what puts runTraining on the .perPixelCE segmentation path.

240 is the native BraTS in-plane size and 240 = 16*15, so a depth-4 UNet's four halvings divide evenly — no resize is needed anywhere in the pipeline. Augmentation is identity for now.

Box detection at 224/7×7 (the YOLOv1 record). 224×224 RGB images; the "labels" buffer carries the 30×7×7 float32 target tensor concatenated with the 7×7 float32 per-cell objectness mask (6076 bytes per record). The runTraining dispatch splits this into target + mask before calling trainStepAdamF32Yolov1. See scripts/datasets/preprocess_visdrone.py for the on-disk format.

BraTS at 224×224 — bratsIO with one number changed. Everything that makes BraTS BraTS (4 modalities, the mask-as-labels buffer that selects the .perPixelCE path, MSD's label numbering for the regions) is identical, so the region table is shared rather than restated: getting WT/TC/ET out of sync between two copies of the same dataset is exactly the silent bug bratsIO's own comment warns about.

BraTS at 224×224 in 2.5D — brats224IO with the channel count changed: 4·(2·ctx+1) channels per record, the centre slice's mask as the label. Same regions, same label numbering, same split; the loader takes the channel count and checks the file length against it, so a 2D file read as 2.5D (or the reverse) fails at load rather than misaligning every mask after the first.

The named class-unions ds reports Dice for (DatasetIO.segRegions; empty for most).

Adam or SGD+momentum (TrainConfig.optimizer), 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.

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. Requires compileVmfbs to have produced (or cached) the _fwd_eval.vmfb already.

End-to-end: compile all three vmfbs, load the train-step session, and run the training loop on the chosen dataset. The entry point the reference trainers call.

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.