Documentation

The constructor's name when this emitter has no lowering for it, else none. The forward and train walks have no arm for these: they write an UNSUPPORTED comment and leave the activation and the parameter list as they were, so the layer would silently become the identity (or a shape error downstream). layerNorm and convNextStem are lowered only by the JAX emitter.

What this emitter cannot render faithfully, or none:

• a layer with no emitter (noEmitter);
• pad := .valid on a .conv2d/.convBn, which this emitter would pad to keep its spatial size;
• a bare .convBn with NetSpec.convBnAct other than ReLU, which this emitter would apply ReLU to anyway.

The JAX emitter implements all three; no spec on this path uses them, so they are refused rather than implemented.

Refuse a spec this emitter cannot render faithfully (unsupported).

What the generate* entry points return for an unsupported spec: a module that fails to parse, whose first line says why. They are pure, so they cannot throw, and a panic! in compiled code only prints and continues with an empty string.

The NCHW input channels when the first layer takes an image (a conv or a conv block); none for a flat input.

Total flat input size: ic · imageH · imageW for an image input (inputChannels); otherwise the first dense layer's fanIn, a token model's sequence (or one-hot) width, or imageH · imageW.

Render tensor type: tensor<128x1x28x28xf32>.

Sanitize a string for use as an MLIR identifier: lower-case, every non-alphanumeric character replaced by _. NetSpec.sanitizedName applies it to a spec's name; the trainer applies it to a spec's buildTag.

The spec's name as an MLIR identifier: the base of every module name the codegen emits (@<sanitizedName>, @<sanitizedName>_eval, jit_<sanitizedName>_train_step, …).

Compute SAME padding. Returns (padH0, padH1, padW0, padW1).

Standard conv dimension numbers block (NCHW/OIHW).

Standard conv attribute block (NCHW/OIHW), symmetric padding, stride 1.

General conv attribute block with explicit asymmetric padding and optional dilations.

Emit an activation forward as inline StableHLO. Takes a uniqueness tag (e.g. s!"_mb{p}"), the input SSA, and the tensor shape. Returns (code, output_ssa). All five activation kinds handled inline so callers don't need to forward-reference helpers further down in this file.

Emit a conv2d + bias + activation block. Returns (code, newSSA, newShape). Assumes NCHW input [b, ic, h, w], produces [b, oc, h, w] for SAME padding, stride 1.

Emit convBn: conv (possibly strided) + instance norm + optional ReLU. Params: W{pidx} (kernel), g{pidx} (gamma/scale), bt{pidx} (beta/shift). Instance norm: per-sample, per-channel spatial statistics.

Emit depthwise convBn: depthwise conv + BN + optional ReLU6. Weight shape: (channels, 1, kSize, kSize), feature_group_count = channels. Output channels = input channels. If noAct := true, emits only convBn (no activation), returning pre-activation so the caller can attach its own (e.g., Swish for MBConv).

Emit a raw depthwise conv (no BN, no activation). Used by ConvNeXt blocks where the DW output goes straight into LayerNorm. Weight shape (channels, 1, kSize, kSize), bias [channels].

Emit an inverted residual block stage for inference. Returns (code, newSSA, newShape, newPidx).

Emit Squeeze-and-Excitation block. Input: x : (B, mid, H, W). Output: x * σ(W_exp · swish(W_red · GAP(x) + b_red) + b_exp). Uses dense (dot_general) in place of 1×1 convs on (B, mid) / (B, seMid) tensors. Params: reduce at pRed (W: [mid, seMid], b: [seMid]), expand at pExp (W: [seMid, mid], b: [mid]). NOTE: weight shapes in MLIR params are stored [seMid, mid, 1, 1] / [mid, seMid, 1, 1] (conv style) — we reshape them to dense layout here. Returns (code, outSSA, shape) where shape == input shape.

Emit an MBConv (EfficientNet) block stage for inference. Same structure as InvertedResidual but:

• uses the layer's kSize (not hardcoded 3)
• uses Swish activation (x * sigmoid(x)) instead of ReLU6
• optional Squeeze-and-Excitation between depthwise and project. Returns (code, newSSA, newShape, newPidx).

Emit a Fused-MBConv (EfficientNet V2) block stage for inference. Unlike regular MBConv, fused-MBConv replaces the 1×1 expand + k×k depthwise pair with a single regular k×k convolution. Structure per block:

1. Fused expand: k×k regular convBn (blockIc → mid) + Swish
2. Optional SE
3. Project: 1×1 convBn (mid → oc), linear (SKIPPED if expand == 1)
4. Skip if blockIc == oc && stride == 1 Returns (code, newSSA, newShape, newPidx).

Emit h-swish forward: x * ReLU6(x + 3) / 6 Returns (code, output_ssa). Caller provides input_ssa and shape.

Emit h-sigmoid forward: ReLU6(x + 3) / 6

Emit ReLU forward (plain, not ReLU6)

Emit an MBConvV3 (MobileNetV3) block for inference. Single block.

• expandCh is the absolute mid channel count.
• If expandCh == ic, skip the expand 1×1 convBn.
• Activation: h-swish if useHSwish, otherwise plain ReLU.
• SE: uses ReLU on reduce, h-sigmoid on gate.
• Skip connection iff stride == 1 && ic == oc.

Emit a Universal Inverted Bottleneck (UIB) block for inference. Single block. MobileNet V4 building block. Structure:

1. Optional preDW (depthwise convBn, NO activation — timm's dw_start); carries the stride only when there is no postDW
2. Expand 1×1 convBn + plain ReLU
3. Optional postDW (depthwise convBn + plain ReLU — timm's dw_mid); carries the stride
4. Project 1×1 convBn, NO activation
5. Skip if stride == 1 && ic == oc All activations are plain ReLU (NOT ReLU6, NOT Swish, NOT h-swish). Stride placement and the activation-free preDW follow timm 1.0.28's UniversalInvertedResidual (scripts/parity/mnv4_timm_parity.py).

Emit global average pool: (b, c, h, w) -> (b, c).

Emit max pool. Handles kSize != stride (e.g., 3/2 in ResNet).

Emit flatten: (b, c, h, w) -> (b, chw).

Channel-concatenate two NCHW tensors A : (N, Ca, H, W) and B : (N, Cb, H, W) along the channel axis (dim 1). Output: (N, Ca + Cb, H, W). Used by UNet decoder skip-connection fusion (emitUnetUp will call this). The two inputs must share (N, H, W).

1-D bilinear-resampling weight matrix Wy : (outLen × inLen) for integer upsampling factor scale, where outLen = inLen * scale. Half-pixel centers, align_corners = false (PyTorch / JAX default): output index i samples input fractional coordinate (i + 0.5) / scale - 0.5, with edge values clamped (so the first and last output rows replicate the corresponding input rows). Each row has at most two nonzero entries.

Serialize a 2-D float matrix as MLIR dense<[[...]]> payload (no dense<> wrapper, no type — caller adds those).

Emit bilinear upsample: (b, c, h, w) → (b, c, hscale, wscale). Factorizes as two dot_generals with precomputed weight matrices: Y = Wy · X (along H) → shape (b, c, w, oH) T = transpose Y to (b, c, oH, w) Z = T · Wxᵀ (along W) → shape (b, c, oH, oW) Wy and Wx are emitted as stablehlo.constant dense<...>. The matmul factorization keeps the cost at O(b·c·(h·w·oH + oH·w·oW)) rather than naïve gather + 4-corner blend, and reuses the existing dot_general lowering paths.

Emit patch embedding forward. Input (B, ic, H, W) → (B, N+1, dim). Three params: W (dim, ic, p, p), b (dim,), cls (dim,), pos (N+1, dim). Returns (code, outSSA, outShape).

Emit LayerNorm forward for a (B, N, D) input, normalizing across D. Returns (code, outSSA, saved records: normSSA, istdSSA, meanSSA). gamma, beta are (D,).

LayerNorm over the channel axis of an NCHW tensor [b, c, h, w]. For each (b, h, w) location, normalize across the c axis with γ, β broadcast from [c].

Implementation note: rather than reducing directly over axis 1 of the NCHW tensor (a [B, C, H, W] shape with the reduce-axis sandwiched in the middle, which IREE's GPU codegen failed to distribute on shapes like Bx(HW)xC), we transpose to NHWC and reshape to [B, HW, C] so the reduction lands on the inner axis — the easy case. The output is transposed back to NCHW so callers see no shape change. Saved norm/istd/mean stay in the inner-axis layout ([B, HW, C] and [B, HW]), which is what the matching backward consumes.

Returns (code, outSSA, normSSA, istdSSA, meanSSA).

Standalone BatchNorm over an NCHW tensor — γ/β per channel, mean/var aggregated over batch+spatial. Uses the same %cbn_* SSA naming as emitConvBn so the existing BN-stats output machinery (NetSpec.bnLayers, train-step BN-stat returns, eval-forward fixed-BN params) all just work when this primitive is registered with the right pidx.

Reductions are over [0, 2, 3] — the well-supported pattern that convBn already uses (and that IREE distributes cleanly on ROCm, unlike LN-NCHW's channel-axis reduction).

Returns (code, outSSA, normSSA, meanBcSSA, istdBcSSA). meanBcSSA and istdBcSSA are the broadcast-to-NCHW versions used by the matching backward.

Emit GELU forward (tanh form, exact) for a tensor of given shape. Returns (code, outSSA, tanhSSA) where tanhSSA is saved for backward.

Emit dense layer for rank-3 (B, N, D) input → (B, N, fanOut). W shape (D, fanOut), b shape (fanOut,).

pairTile forward, shared by the eval and train emitters: the flat host input [B, 2·L·D] (residue block i, then block j) reshaped to [B, 2L, D] and split; the i block times W, the j block times Wj (no bias: the convBn after it has β); the two [B, L, C] results broadcast against each other and summed into [B, L, L, C], then transposed to NCHW [B, C, L, L]. With p > 0 pair planes the host row is [2·L·D | p·L·L]: the feature blocks are its head, and the tail, reshaped to [B, p, L, L], is concatenated onto the map as its last p channels. Returns (code, outSSA, uSSA, vSSA), the two block inputs kept for the backward's weight gradients.

Emit the FlashAttention forward SDPA core: softmax(Q·Kᵀ/√dh + mask)·V via a tiled online-softmax stablehlo.while over ⌈n/bk⌉ key blocks, so the only O(n·bk) tensor live per iteration is one score block [b,h,n,bk] — the full [b,h,n,n] matrix is never materialized. Inputs q/k/v : [b,h,n,dh]; returns (code, oSSA [b,h,n,dh], lseSSA [b,h,n]). lse = m + log(l) is the per-row logsumexp the backward needs. Requires n % bk == 0. Validated end-to-end against dense attention (jax/demos/flash_attention_ref.py + an IREE compile+run probe, ~1e-7 fp32).

Emit the FlashAttention backward: given q/k/v [b,h,n,dh], the upstream grad dO [b,h,n,dh], the forward output O, and the per-row logsumexp lse [b,h,n], produce (dQ, dK, dV) — all [b,h,n,dh] — via a recompute stablehlo.while over key blocks. No [b,h,n,n] matrix is stored: each block recomputes Pij = exp(Sij − lse) and uses D = rowsum(dO⊙O) as the softmax-Jacobian term. dQ accumulates across ALL blocks; dK/dV are written once per (disjoint) key block via dynamic_update_slice. Matches flash_attention_ref.py:flash_backward (validated ~1e-15).

Block size for FlashAttention: a divisor of n ≤ 128 (fewer blocks for short sequences, real memory win for long ones).

Emit the RoPE cos/sin tables [n, dh] in-graph (scale-independent — no giant constant): θ_i = base^(-2i/dh), freqs[m,i] = m·θ_i, tiled to dh. Returns (code, cosSSA, sinSSA). See jax/demos/rope_ref.py.

Apply RoPE to xSSA : [b,heads,n,dh] given the [n,dh] cos/sin tables: x⊙cos + rotate_half(x)⊙sin, rotate_half([a,b])=[-b,a]. With negSin (the backward/VJP), sin is negated (rotation by −θ). Returns (code, outSSA).

Emit multi-head self-attention forward. Input x (B, N, D). Params: Wq, bq, Wk, bk, Wv, bv, Wo, bo. Returns (code, outSSA, Q, K, V, softmax-or-Lse, preProj, sdpaOut SSAs).

Emit a transformer encoder block forward pass. Layout: LN1 → MHSA → + → LN2 → fc1 → GELU → fc2 → + Uses pidx starting at startP and 8 param pairs per block: LN1(g,b), Wq,bq, Wk,bk, Wv,bv, Wo,bo, LN2(g,b), Wfc1,bfc1, Wfc2,bfc2. Returns (code, outSSA, newPidx).

Emit a residual block stage: nBlocks basic blocks, first may downsample. Returns (code, newSSA, newShape, newPidx).

Emit a bottleneck block stage: nBlocks bottleneck blocks (1x1→3x3→1x1). mid = oc / 4 is the bottleneck channels. First block may downsample.

Emit a dense layer. Assumes (batch, fanIn) input.

Emit a ConvNeXt stage (forward inference): nBlocks repeats of DW7×7-raw → norm(channel-axis) → 1×1 expand (×4) → activation → 1×1 project → LayerScale → residual add. Five pidx slots per block, in the order: DW (W, b), norm (γ, β), expand (W, b), project (W, b), LayerScale (γ). norm selects LN (default, paper) or BN (running stats fed in eval). Caller must already have emitted any preceding downsample.

Emit a ConvNeXt downsample (forward inference): norm(channel-axis) on ic-channel input, then a 2×2 stride-2 valid-pad conv ic → oc. Two pidx slots: norm (γ, β), conv (W, b). norm selects LN/BN.

The per-image timestep is carried by the last channel of the DDPM network input (prependTChannel fills it with t/Tmax, constant over the plane). Given the flat input %x_flat : [B, inDim] and the channel layout, the timestep column offset is (cIn-1)·H·W.

Forward emit for .timeCondAdd channels nFreq, shared by the eval and train-step forward walks. Reads the per-image timestep scalar s = t/Tmax from %x_flat[:, tOff], builds a 2·nFreq sin/cos embedding in-graph (frequencies π·2^k), projects it through the learned dense pW : [2·nFreq, C] + pB : [C], and adds the result (broadcast over H,W) onto featSSA : [B, C, H, W]. Returns (code, outSSA, embSSA) — embSSA (the [B, 2·nFreq] sin/cos features) is stored by the train forward for the backward.

Token-embedding forward via a true stablehlo.gather (Part II Option 2). The [B, T] f32 ids index the [V, D] table directly: emb[b,t,:] = W[ids[b,t], :], producing %tpe_emb{pos} [B, T, D] — no [B, T, V] one-hot. f32 ids convert to i32 exactly (vocab « 2²⁴). Literal attribute braces are concatenated (Lean s-strings reserve {).

1×1 conv (channel mix) forward: einsum('oi,bihw->bohw', W, x). x [b,ic,h,w], W [oc,ic] → [b,oc,h,w]. dot_general contracting the input channel dim, then transpose the trailing out-channel back to NCHW.

1×1 conv input-VJP: einsum('oi,bohw->bihw', W, dP) → dX [b,ic,h,w]. Contract the out-channel dim of dP against W's out dim.

1×1 conv weight-VJP: einsum('bohw,bihw->oi', dP, x) → dW [oc,ic]. Contract the batch + spatial dims jointly.

Standalone VJP of emitBilinearUpsample ([b,c,h,w] → [b,c,h·s,w·s]): the transpose matmul-pair. Mirrors the reverse-record .bilinearUpsample backward exactly, but as a reusable emitter so the FPN DAG can route an upsample cotangent up the pyramid. inShape is the (small) pre-upsample shape; gradSSA carries the (large) post-upsample cotangent.

FPN top-down neck forward. Backbone taps C3/C4/C5 at strides 8/16/32 (grids 4·g5 / 2·g5 / g5), lateral 1×1 convs W3/W4/W5 : [oc,c_n]:

P5 = conv1x1(C5,W5)                  [b,oc,g5,g5]
P4 = conv1x1(C4,W4) + upsample2(P5)  [b,oc,2·g5,2·g5]
P3 = conv1x1(C3,W3) + upsample2(P4)  [b,oc,4·g5,4·g5]

Returns the code and the three pyramid SSAs (P3,P4,P5).

FPN neck VJP. Each Pn cotangent routes to its lateral conv AND up the pyramid (accumulating: dP4_tot = dP4 + up^T(dP3), then dP5_tot = dP5 + up^T(dP4_tot)). Returns the code, the backbone-tap cotangents (dC3,dC4,dC5), and the lateral-weight gradients (dW3,dW4,dW5). Mirrors fpn_grad in scripts/probes/fpn_neck_check.py.

One k×k stride-1 same-pad conv + per-channel bias + ReLU, as a standalone prefix-parameterized emitter. Channel-preserving (c → c), which is all the RetinaNet head tower needs. The math is lifted verbatim from the FD-verified .conv2d forward/backward in the reverse-record walk; this just makes it callable from inside a single layer's emitter, the same way emitConv1x1Fwd/BwdDC/BwdDW serve the neck.

Returns the code, the post-ReLU output SSA, and the PRE-activation SSA (the backward needs it for the ReLU mask).

VJP of emitTowerConvFwd: ReLU mask → dW, db, dX. preActSSA is the pre-ReLU SSA the forward returned; xSSA is that layer's input. Returns the code, dX, dW, db. Same three-step structure as the .conv2d backward (dW by the transpose trick, dX by the reversed+transposed kernel).

RetinaNet head tower for ONE pyramid level: depth × (3×3 conv → +bias → ReLU), channel-preserving. Towers are per-scale, not shared across P3/P4/P5, which sidesteps the shared-weight gradient accumulation (and, if norm is ever added, the per-level-statistics problem) that sharing would introduce.

No normalization: plain conv + bias + ReLU, as in the original RetinaNet head. That keeps the tower out of the BN running-stats plumbing entirely.

Returns the code, the output SSA, and the per-layer (input, preAct) SSA pairs the backward needs, in forward order.

VJP of emitFpnTowerFwd: walk the layers in REVERSE, threading the input cotangent back to the tower's input. Returns the code, dX, and the per-layer (dW, db) pairs in FORWARD order (so they bind to %W{base+i} directly).

Param-index layout for a .fpnDetect layer based at base, as SSA names with prefix pfx ("%W", "%m_W", "%v_W", "%d_W", …). ONE definition of the ordering, so the forward, backward, signature, optimizer and paramShapes cannot drift apart:

base + 0 .. base + 2                  Wn3, Wn4, Wn5           (neck laterals)
base + 3 .. base + 3 + 6·depth − 1    tower (W, b) × depth, P3 → P4 → P5
headBase + 0 .. headBase + 2          Wh3, Wh4, Wh5           (head convs)
headBase + 3 .. headBase + 5          bh3, bh4, bh5           (head biases, LAST)

where headBase = base + 3 + 6·depth. Returns the per-level tower weight and bias name lists; depth = 0 gives three empty lists.

Index of the first head param (Wh3) for a .fpnDetect based at base.

Total param count owned by a .fpnDetect layer: 9 + 6·depth.

FPN detector head+neck forward: neck (top-down merge) → per-scale 1×1-conv heads → flatten + concat into a single [B, Ntot] output. Heads are a single 1×1 conv Pn[oc] → [A·15] plus a per-channel bias (no hidden tower — the minimal detection head). Params: neck laterals wn3/wn4/wn5 : [oc,c_s], head convs wh3/wh4/wh5 : [A·15, oc], head biases bh3/bh4/bh5 : [A·15]. Concat order P3,P4,P5 matches emitMultiScaleYoloLoss's split. Returns the code, the concat SSA, and the pyramid SSAs (P3,P4,P5 — needed for the head weight VJP).

The bias is what carries the RetinaNet prior init (applyDetPriorBias). A biasless 1×1 conv has to manufacture the constant background offset out of its weights, which is measurably what it spends them on: on a biasless run's validation set every objectness logit sits in [−2.7, −1.2]. Zero-init biases reproduce the biasless net exactly.

FPN detector backward: un-concat the [B, Ntot] loss grad into the 3 head grads → 1×1-conv head VJP (dPₙ + dWhₙ + dbhₙ) → neck VJP (dC3/dC4/dC5 + dWnₙ). Returns the code, the backbone-tap cotangents (dC3,dC4,dC5), and the 9 param grads [dWn3,dWn4,dWn5, dWh3,dWh4,dWh5, dbh3,dbh4,dbh5] (order matches the forward's param args). The bias VJP is the plain sum of the head cotangent over batch + spatial (a bias added to every cell receives every cell's grad). All from FD-verified conv1x1 / neck emitters.

Emit the @forward function body walking layers in order. When stopAtGAP is true (GradCAM mode), the walk halts at the first .globalAvgPool it encounters: instead of emitting GAP+dense, we reshape the pre-GAP activation [B, C, H, W] to flat [B, C*H*W] and return that. CAM math is then done in Lean against the dense head's weights.

Split a layer's ParamSlots into signature groups: each W opens a group, and a group's tensors share one parameter index (%W3, %g3, %bt3).

One layer's forward-signature parameter arguments starting at index p, ,\n %W3: …, %g3: …, %bt3: … per group — read off Layer.paramSlots — and the next index. .fpnDetect keeps its own line layout (laterals, tower pairs, head weights, head biases); .tokenPositionEmbed puts both tables on one line.

The forward signature's parameter list and output shape: %x, the parameter arguments (fwdLayerArgs), and the activation shape walked through every layer. Shared by emitForwardSig and emitForwardEvalSig.

Emit function signature: interleaved (W, b) per param-bearing layer.

Generate a StableHLO MLIR module for forward-only inference.

BN-layer (pidx, oc) pairs, one per convBn layer (including those inside residual/bottleneck blocks), in the order the codegen threads their running stats. ViT-style transformer specs return an empty array.

Emit function signature for forward_eval: same as forward but with bn_mean/bn_var params appended for each BN layer.

Generate a StableHLO MLIR module for eval (fixed BN using running mean/var).

Shape walker: returns (C, H, W) of the activation immediately before the LAST .globalAvgPool in the spec — i.e. the last conv feature map, the input to the GAP+dense head. Returns none for specs without a globalAvgPool (those are not CAM-eligible). Mirrors the shape transitions in emitForwardEvalSig; kept focused (no param bookkeeping) so it stays readable.

Forward-with-capture signature: same params as emitForwardEvalSig but returns the flat pre-GAP feature map [batch, C*H*W] instead of the logits [batch, NC].

Generate a StableHLO MLIR module for forward-with-capture (CAM mode). Module name <sanitized>_cam, function @forward_cam. Walks the network up through the last conv feature map, reshapes it to flat [batch, C*H*W], and returns. The Lean side recomputes logits and the CAM heatmap from the dense-head weights.

Emit convBn forward for train step. Records intermediates for backward. Returns (code, FwdRec).

Emit depthwise convBn forward for train step. Records intermediates for backward. Weight shape: (channels, 1, kSize, kSize), feature_group_count = channels. Activation: ReLU6 by default, Swish if useSwish := true, or h-swish if useHSwish := true.

Emit convBn forward for train step with ReLU6 (used for inverted residual expand). Same as emitConvBnTrain but with ReLU6 activation instead of ReLU.

Emit convBn forward for train step with Swish (x * sigmoid(x)) activation. Used by EfficientNet's MBConv.

Emit convBn forward for train step with h-swish activation: h_swish(x) = x * ReLU6(x + 3) / 6. Used by MobileNetV3.

Emit convBn backward: inst norm VJP + conv backward. Returns (code, gradient SSA name, gradient shape).

Emit depthwise convBn backward: BN VJP + depthwise conv backward. Weight shape: (channels, 1, kSize, kSize), feature_group_count = channels. Returns (code, gradient SSA name, dW SSA name).

Emit Adam update for one param. m_new = β1*m + (1-β1)grad v_new = β2v + (1-β2)*grad² m_hat = m_new / bc1 (bc1 = 1 - β1^t, precomputed) v_hat = v_new / bc2 (bc2 = 1 - β2^t, precomputed) w_new = w - lr * m_hat / (sqrt(v_hat) + ε)

• optional decoupled weight decay: w_new = w_new - wdlrw

Emit SGD+momentum update for one param. v_new = mu * v + grad w_new = w - lr * v_new

• optional decoupled weight decay: w_new = w_new - wdlrw

Emit the per-pixel softmax-CE block for segmentation. Logits are (B, NC, H, W) (curShape at the point of call), labels are (B, H, W) int32 (passed in as %y_seg by the caller).

Forward: loss = -mean over (n,h,w) of log softmax(logits)[n, label[n,h,w], h, w] Backward: %d_logits_seg = (softmax - onehot(labels)) / (B · H · W)

Returns the MLIR text. The caller is responsible for setting gradSSA := "%d_logits_seg" and gradShape := [B, NC, H, W] after appending this block.

lossOut/gradOut name the two results. They default to the plain %loss / %d_logits_seg the seg path consumes, and are overridden only by .diceCE, which emits this block and the Dice block side by side under private names and sums them.

Soft-Dice loss + d_logits for the segmentation path.

Forward, per class c, reducing over the whole batch (dims 0,2,3):

I_c = Σ p_c·y_c P_c = Σ p_c Y_c = Σ y_c D_c = (2·I_c + ε) / (P_c + Y_c + ε) loss = 1 - (1/NC)·Σ_c D_c

ε (smooth) both regularizes the ratio and makes an absent-and-correctly- -unpredicted class score D_c = ε/ε = 1 (zero loss) rather than 0/0.

Backward. D_c depends only on channel c, so dL/dp is elementwise with two per-class broadcast scalars:

∂D_c/∂p_{c,k} = [2·y_{c,k}·den_c - num_c] / den_c² = A_c·y_{c,k} - B_c where A_c = 2/den_c, B_c = num_c/den_c² g_{c,k} := ∂L/∂p_{c,k} = (1/NC)·(B_c - A_c·y_{c,k})

Then chain through the channel softmax. Unlike CE — whose p - y seed is what's left after the softmax Jacobian cancels the log — Dice gets no such cancellation, so the Jacobian-vector product is explicit:

dz_i = p_i·(g_i - Σ_j g_j·p_j)

which is elementwise ops plus one channel reduce. No new primitive.

Labels are one-hot'd raw here: label smoothing is a CE notion (it softens a log-likelihood target) and has no meaning for a set-overlap ratio, so .diceCE applies smoothing to its CE half only.

Per-pixel focal CE (Lin et al.) + d_logits for the segmentation path.

Forward, per pixel k, with p_t the softmax probability of the true class:

FL_k = -(1 - p_t)^γ · log p_t loss = (1/N)·Σ_k FL_k

Backward. Write omp := 1 - p_t. Differentiating through p_t and then the channel softmax (∂p_t/∂z_c = p_t·(δ_tc - p_c)) collapses to one per-pixel scalar times the familiar CE shape:

A_k := γ·p_t·omp^(γ-1)·log p_t - omp^γ ∂FL_k/∂z_c = A_k · (y_c - p_c)

Two readings of A that are the whole reason this loss is in the demo:

• γ = 0 ⇒ A = -1 ⇒ dz = (p - y)/N, i.e. exactly perPixelCE. That is a free structural test, and seg_loss_probe_check.py runs it.
• p_t → 0 (the collapsed class) ⇒ p_t·log p_t → 0 and omp^γ → 1, so A → -1 and the gradient tends to CE's, not to zero. Focal does not rescue a collapsed class by amplifying it — contrast Dice, whose p_i factor sends the gradient to zero exactly there. It works, when it works, by driving A → 0 at p_t → 1 and defunding the easy majority.

omp^γ is emitted as exp(γ·log(max(omp, ε))) rather than stablehlo.power, matching the YOLOv1 objectness path's idiom. The ε clamp only engages at p_t > 1 - ε (a saturated softmax), where omp^γ is ~0 and the pixel contributes nothing anyway.

Unlike that YOLO path, the weight here is not detached: A carries the derivative of the (1-p_t)^γ factor itself. Detaching is a defensible approximation but yields a d_logits that is not the gradient of the loss being reported, which FD would (correctly) reject.

Emit the seg loss block selected by segLoss, leaving %loss and %d_logits_seg defined either way. .diceCE emits both blocks under private names and sums them; the duplicated softmax prologue is identical SSA on identical inputs, so CSE folds it.

DIoU box-loss FORWARD block (the IoU-family replacement for the fragile √-MSE box regression). pred/tgt are [B,4,gH,gW]: pred channels are the raw head outputs (tx,ty,tw,th), tgt channels are the YOLOv1 target layout (cell-offset x,y then w_rel,h_rel). mask is [B,gH,gW] (1 on object cells). Positive box parameterization by construction — cx=(j+σ(tx))/gW, cy=(i+σ(ty))/gH, w=exp(tw), h=exp(th) — so boxes are always valid (fixes the negative-w/h gradient death). Emits loss = Σ_cells mask·(1 - DIoU) with DIoU = IoU − ρ²(centers)/c²(enclosing). Requires %zf in scope. Numeric spec: scripts/probes/diou_grad_check.py; the emitted forward is checked against numpy in scripts/probes/diou_probe_check.py.

DIoU box-loss BACKWARD block. Assumes emitDiouForward ran just before (reuses its %{pfx}_* SSA values). Emits {gradOut} = ∂loss/∂pred [B,4,gH,gW] (channels d/dtx, d/dty, d/dtw, d/dth), the hand-derived VJP of Σ mask·(1 - DIoU). Analytic spec: diou_loss_grad in scripts/ diou_grad_check.py; FD-verified against the emitted forward in scripts/probes/diou_probe_check.py. Every piecewise min/max sub-gradient is an indicator (compare→select).

Anchor YOLOv3-style loss. A anchors/cell, each slot is [tx,ty,tw,th, obj, cls(NC=10)]; pred/tgt are [B, A·15, gH, gW], mask is [B, A, gH, gW] (per-anchor objectness assignment). Per anchor a: the FD-verified DIoU box loss at pfx=a{a} with prior anchors[a] (box_a = anchor_a·exp(pred)), focal-BCE objectness, and mask-weighted softmax-CE class. Loss = (1/B)·Σ_a [λ_box·Σmask_a(1-DIoU_a) + Σ focalBCE_a + Σ mask_a·CE_a]; the gradient is assembled per anchor to [B,15,gH,gW] and concatenated to [B, A·15, gH, gW]. Requires %zf,%neginf in scope. FD-verified in scripts/probes/anchor_loss_probe_check.py.

FPN multi-scale YOLO loss over the CONCATENATED head output (bites 4+6). The 3 detection heads each flatten [B, A·15, gn, gn] → [B, A·15·gn²] and concat into logitsSSA [B, Ntot]. This splits it back per scale, runs the FD-verified emitAnchorYoloLoss on each (a distinct s{i} tag per scale so the loss-level SSAs don't collide), sums the per-scale losses, and re-concats the per-scale grads back to [B, Ntot] — the single-output adapter that keeps emitTrainStepBody single-output. Scales are independent (no cross-scale terms), so d(loss)/d(logits) = concat(flat gradₛ). Requires %zf in scope. scales = per-scale (grid, anchors); tgtSSAs the matching target blocks.

The train step's constants: %zf / %neginf, the optimizer's (Adam or momentum), and the weight decay.

The loss on the logits logitsSSA and its gradient seed — FPN multi-scale YOLO, YOLOv1, DDPM, per-pixel segmentation, or softmax cross-entropy (optionally focal / soft-label). Returns the text, the seed's name and its shape, which the backward walk starts from.

The optimizer updates, one per Layer.paramSlots tensor, then the return of the updated θ / m / v, the loss and the BN statistics.

The train step's forward walk: one block of StableHLO per layer, from the input reshape to the logits. Returns the text, the logits' SSA name and shape, and one FwdRec per layer — the names and shapes emitTrainBackward reads back.

The train step's backward walk: the FwdRecs in reverse, each layer's VJP from the cotangent gradSSA₀ : gradShape₀ the loss seeds, writing %d_W{p} / %d_b{p} (and the other parameter gradients) that emitOptimizerUpdates reads. Returns the text and the cotangent the walk ends on: the input's, [B, ic, H, W] after a first conv, [B, fanIn] after a first dense.

emitTrainBackwardTo's text alone: the train step discards the cotangent the walk ends on.

Emit the full train step: constants, emitTrainForward, emitTrainLoss, emitTrainBackward, emitOptimizerUpdates.

The train step's parameter arguments with prefix pfx ("" θ, "m_", "v_"), one line per group, and the parameter types in order — read off Layer.paramSlots. Two layers keep their own line shapes: .fpnDetect numbers every tensor, one per line, and .tokenPositionEmbed puts its embedding and positional table on one line.

Emit the train_step function signature.

Generate a full train_step MLIR module with VJPs. When useSoftLabels is true, the function takes a %y_soft : [B, NC] f32 tensor (mixup/ cutmix-style) instead of an int32 %y : [B] label vector; label smoothing is then expected to be applied by the caller.

The layers generateVjp takes: the plain conv / dense stacks whose backward walk ends on the input's cotangent (each .conv2d emits its %d_x, each .dense its %d_in, .flatten the inverse reshape), with ReLU or identity activations (the walk's conv / dense backward masks for ReLU and passes any other activation's cotangent through untouched). A block (BN, residual, attention, …) can end its walk on a skip or token gradient instead, so it is refused rather than trusted.

The trained tensors in trainSigArgs order, each with its index-suffixed name (W0, b0, …): the names emitTrainBackward writes as %d_<name>.

The vector–Jacobian product of the network, as its own module:

@main(θ…, %x_flat : [B, inDim], %gbar : [B, …out]) -> (∂θ…, %d_x_flat : [B, inDim])

with ∂θ = (∂net/∂θ)ᵀ ḡ and ∂x = (∂net/∂x)ᵀ ḡ, summed over the batch (no 1/B, no loss). It is the train step's own forward and backward walk (emitTrainForward, emitTrainBackwardTo) seeded with the caller's cotangent instead of a loss's, with no optimizer tail: a network composed with others — MuZero's representation, dynamics, prediction and reward nets — takes its cotangent from the next one and hands ∂x back. Gated by scripts/probes/vjp_gate.py: ∂θ against the SGD train step at lr 1, ∂x against central finite differences and jax.vjp. Plain conv / dense stacks only (vjpUnsupported).

One Adam step on a given gradient, as its own module:

@main(θ…, m…, v…, ∂θ…, %lr, %t) -> (θ'…, m'…, v'…)

the train step's own emitAdamUpdate per tensor (β = 0.9 / 0.999, ε = 1e-8, bias-corrected at step t, decoupled weight decay on the ParamSlot.decay tensors when weightDecay > 0). For a gradient the caller assembled: a net applied K times in one unroll takes the sum of its K generateVjp gradients.

Standalone segmentation-loss module — exercises emitSegLossBlock in isolation so the loss and its gradient can be checked against finite differences: @main(logits : [B,NC,H,W], y : [B,H,W] i32) -> (loss, d_logits).

This is the FD harness for .dice / .diceCE, whose gradient (unlike CE's p - y) is a real softmax Jacobian-vector product and therefore worth checking numerically rather than by inspection. Driven by scripts/probes/seg_loss_probe_check.py; CPU is fine, no GPU needed.

Standalone DIoU box-loss probe: @main(pred,tgt : [B,4,gH,gW], mask : [B,gH,gW]) -> (loss, d_pred). Compiled + run on CPU by scripts/probes/diou_probe_check.py — the emitted forward is checked against numpy and the emitted backward against central finite differences, before wiring into the YOLOv1 train step.

Standalone anchor-YOLO-loss probe: @main(pred,tgt : [B,A·15,gH,gW], mask : [B,A,gH,gW]) -> (loss, d_pred). Deterministic test anchors (w=0.02+0.03·i, h=0.03+0.04·i), γ=2, λ_box=5 — mirrored in scripts/probes/anchor_loss_probe_check.py, which checks the emitted forward against numpy and the emitted backward against finite differences.

Standalone FPN-neck probe: @main(C3,C4,C5, W3,W4,W5, dP3,dP4,dP5) -> (P3,P4,P5, dC3,dC4,dC5, dW3,dW4,dW5). Compiled + run on CPU by scripts/probes/fpn_neck_probe_check.py, which checks the emitted forward against the numpy fpn_forward and the emitted backward against the f64-FD-verified fpn_grad oracle. Cotangents are explicit inputs — no scalar-loss reduce — so the probe mirrors the real train-step wiring (head backwards feed dPn tensors straight into the neck VJP).

Standalone FPN multi-scale-loss probe: @main(logits [B,Ntot], T0,T1,...) -> (loss, grad [B,Ntot]), one target block per scale. Exercises the concat split

• per-scale anchor loss + grad re-concat in isolation — conv-free, so it CPU-compiles for FD checking (the conv layers that feed it in the train step are not part of this probe). Deterministic per-scale anchors (w=0.02+0.03·i, h=0.03+0.04·i), matching scripts/probes/fpn_loss_probe_check.py.

Standalone whole-FPN-detector probe: neck + optional RetinaNet tower + 1×1-conv heads + concat + multi-scale loss + full DAG backward. Run at focal γ=0 so the objectness weight is a genuine constant (not just detached) and EVERY input/param grad is exactly finite-differenceable — this validates the DAG backward (all 9 + 6·tower param grads) before the emitTrainStepBody wiring.

@main(C3,C4,C5, Wn3,Wn4,Wn5, [Wt_i,bt_i × 3 levels × tower], Wh3,Wh4,Wh5, bh3,bh4,bh5, T3,T4,T5) -> (loss, dC3,dC4,dC5, <the same param grads in the same order>)

Param args after the neck follow fpnDetectParamShapes exactly, so the probe is also a check that that ordering is what the emitters actually consume. tower = 0 keeps the module conv1×1-only (pure dot_general), which is what lets it CPU-compile; tower > 0 introduces real 3×3 stablehlo.convolution ops, so CPU-compilability there is not guaranteed (see the probe checker). Deterministic per-scale anchors (w=0.02+0.03·i, h=0.03+0.04·i).

Standalone FlashAttention forward module — exercises emitFlashAttnSdpa in isolation for validation against dense attention: @main(Q,K,V : [b,h,n,dh]) -> O. Compiled + run in the flash-attn probe.

Standalone DENSE SDPA forward module @main(Q,K,V) -> O (materializes the full [b,h,n,n] scores) — the memory-comparison baseline for the flash probe. Compile both with --iree-scheduling-dump-statistics-file at increasing n to see dense peak allocation grow O(n²) while flash stays O(n·bk).

Standalone RoPE module @main(x) -> rope(x) (or the VJP with backward), for op-level validation against jax/demos/rope_ref.py.

Standalone FlashAttention forward+backward module: @main(Q,K,V,dO) -> (dQ,dK,dV). Runs the forward emit to get O/Lse, then the backward emit; validated against the dense-attention VJP.