Documentation

Per-parameter tensor shapes for the entire spec, in the same order MlirCodegen.emitTrainStepSig walks them. Used to pack/unpack the params ++ m ++ v ByteArray that goes into and out of every trainStepAdamF32 call.

Packed params ++ m ++ v shape array (3× the param shapes), as int32 LE ByteArray. This is what gets passed to trainStepAdamF32.

Total float count needed to store running BN stats (mean + var per BN layer). 0 for ViT and any non-BN architecture.

BN shapes packed for the FFI: [n_bn_layers, oc0, oc1, ...] as int32 LE. The trainStepAdamF32 FFI uses this to know how many BN-stat outputs to pop after the params/loss.

Param shapes for the eval forward pass: regular params followed by one [oc] mean and [oc] var per BN layer. ViT-style specs collapse this to just paramShapes because there are no BN layers.

Packed evalShapes as int32 LE for the eval forward FFI.

Packed input-tensor shape for a flat-image batch (NCHW collapsed to [batch, channels*H*W]). Channel count comes from the first conv-style layer; defaults to 1 for pure-MLP specs.

The eval forward function name to pass to forwardF32. The codegen emits modules of the form @<sanitized_name>_eval containing func.func @forward_eval, so the FFI call wants the qualified "<sanitized_name>_eval.forward_eval". Computing this from the spec means trainers can never get the spelling wrong by hand.

He-initialize all parameters for a spec, walking layer-by-layer from seed.

Overwrite the head's bias with log priors[c] — RetinaNet's prior-bias init (Lin et al. §3.3, "prior"), for segmentation heads.

heInitParams lays layers out in order and emits each conv's bias directly after its weights, so the head's bias is the final NC floats of the buffer. That is what this patches; it is a no-op on everything else.

What it buys, and why it is focal's other half. A zero-bias head starts at a uniform softmax: every class at 1/NC, background included. The net's first job is therefore to discover the class prior, and on BraTS it does that by walking straight into the trivial predictor — the collapse is decided in the first ~100 steps. A log π_c bias hands it the prior at step 0 instead, so the first gradient step is spent on the actual task.

The quantitative version is in scripts/probes/seg_grad_scorecard.py, whose sweep lands on this exact row. Prior-bias init starts the net at z₀ - z₃ = log(π₀/π₃) = log(0.9746/0.0050) = 5.27 (verified against the emitted checkpoint: softmax(head bias) == π to 2e-09). Its measured balance ratio — rare-class gradient over majority gradient — reads:

One bias vector is worth ~19,000× to focal, at step 0 — and flips it from the worst arm (tied with CE, a literal no-op) to the best (~2× wce). That is the whole content of "focal needs confidence to suppress": this manufactures the confidence up front instead of waiting for training to produce it too late. It is why Lin et al. ship the two together; they are one idea, and reading them as separate tricks is how the pairing gets lost.

It helps every arm — wce's ratio rises 51× too — because starting at the prior is simply a better starting point than uniform. focal is the one that goes from inert to leading.

priors need not be normalized: adding a constant to every logit is a no-op under softmax, so an overall scale on priors shifts every bias by log k and changes nothing. Only the ratios matter — the same scale-invariance that perPixelWeightedCE's reduction enjoys, for the same reason.

Returns a NEW ByteArray.

RetinaNet prior-bias init for the FPN detector head. Sets every objectness logit's bias to −log((1−π)/π) so the head starts predicting sigmoid = π (π ≈ 0.01) on every cell; box and class biases stay at zero.

This is the classifier trick of applyHeadPriorBias transposed to a sigmoid/one-vs-all head, and it is aimed at a measured failure rather than a guess. On a 12-epoch run every objectness logit sat in a narrow negative band, positives barely above negatives: the head had real signal but almost no dynamic range, because a bias-free 1×1 conv has to synthesize the constant background offset out of weights that also have to discriminate. The bias hands it that constant for free, which is the whole point — and per-class mAP is bounded by objectness ranking, so this is the lever that can move it.

Assumes the .fpnDetect layer is last, so its 3 [A·15] biases are the final 3·A·15 floats of the buffer — the same tail-splice applyHeadPriorBias does. Within each [A·15] block, anchor a's objectness is channel a·15 + 4 (emitAnchorYoloLoss slices box at base..base+4, obj at base+4, class at base+5..base+15).

Returns a NEW ByteArray.

Patch the first prefixBytes of initParams with bytes read from pretrainedPath. Used to bootstrap a fresh init from a pretrained backbone — e.g. load R34-Imagenette weights into a YOLOv1 init, keeping the YOLOv1 head's He-init untouched.

prefixBytes is computed by the caller (usually 4 * (spec.totalParams - <last layer fanOut + fanIn*fanOut>) for a spec whose final dense layer differs from the pretrained source).

Returns a NEW ByteArray; initParams is not modified.

Offset-aware sibling of patchInitWithPretrainedPrefix: copy countBytes from pretrainedPath starting at srcOffBytes into initParams starting at dstOffBytes, leaving everything outside that window at its He-init.

The prefix version can only bootstrap weights that are contiguous at the FRONT of the layout, which forces the first layer to match the checkpoint's first layer exactly. That is the wrong constraint the moment the input channel count differs: an R34 trained on 3-channel RGB has a [64,3,7,7] stem, and a 4-modality MRI net needs [64,4,7,7]. The shapes disagree by 3,136 floats, so a prefix patch would land every subsequent weight at the wrong offset — silently, since the sizes still "fit".

With a range, the fresh stem keeps its He-init and the rest of the backbone lands where it belongs: dstOff = stem floats in THIS spec, srcOff = stem floats in the CHECKPOINT, count = backbone floats − srcOff.

Both windows are bounds-checked, because the failure mode this exists to prevent is a silent misalignment rather than a crash.