ConvNeXt-T train step rendered from the verified AST, with declared carve-outs #
The per-example ConvNeXt render. It writes one artifact, the SGD-inline
verified_mlir/convnext_train_step.mlir (convNextTrainStepText), and holds what
ConvNeXtRenderB reuses: the stage tables (CnxDims, cnxTiny/cnxSmall/cnxBase), the
parameter list (cnxAllParams), the forward render convNextFwdText, the shared traversal
convNextBackAll, and the AdamW tail convNextAdamTrainStepText. Every other ConvNeXt
artifact is written by the batched chain in ConvNeXtRenderB, whose Proofs tier is
ConvNeXtFoldGB. This chain stays because the batched traversal has no fused-SGD arm and
ConvNeXtStepTie is stated at this file's bytes. TestConvNeXtFwdBTie checks the two chains
against each other.
The train step is the full [3,3,9,3] ConvNeXt-T (BS=32, 3×224²→10): forward, backward
cotangent chain, and the parameter updates. All 182 parameter gradients are SHlo nodes. The stem
4×4/s4 weight gradient is .convStride4WeightGrad, whose den is
flatConvStride4WeightGradHasVJP; the three 2×2/s2 downsample weight gradients are
.convStridedWeightGrad, whose den is the kernel-generic flatConvStride2WeightGradHasVJP
(the emitter splits odd/even kernels in StableHLO.sWGradGeom).
Hand-written text in the rendered module, outside any SHlo term:
- the function signature, the
%sc/%bscconstants andchLnPrelude; - the cotangent's
÷B(%dy = divide …, %bsc) in the SGD render; - the GAP-backward block (
%dgi…%dgapf), in both the SGD and AdamW renders; - the stem weight's update in the SGD render:
psW − lr·gissgdOf, a hand-writtenmultiply/subtractaround the certified gradient; - the AdamW render's report-only
%loss.
In the SGD render every other parameter is updated by a fused SHlo SGD op (gradient and
θ − lr·g in one node), which denotes θ − lr·(certified per-layer Jacobian · its cotangent)
(ConvNeXtFold). Render is
value-independent (skel erases values), so placeholders + lr:=0 are passed; the emitted
lrStr/epsStr literals carry the real values.
References #
- Liu et al. 2022, A ConvNet for the 2020s (ConvNeXt). https://arxiv.org/abs/2201.03545
The ConvNeXt size — depths and channel dims TOGETHER, in one record.
They are bundled deliberately. A bare Array Nat of depths beside a bare array of dims
admits (depths := S, dims := T), which is not a ConvNeXt of any size but type-checks, renders
and trains. B moves BOTH, so two bare arrays would make the mismatch reachable in the
direction that matters. One record, three instances, and a caller cannot spell a net that does
not exist. Same reasoning as VitDims, arrived at from the opposite direction: ViT bundled to
keep d = heads * hd definitional, this bundles to keep two tables from drifting.
Both tables are Vector Nat 4: ConvNeXt has four stages at every size, and the renders read
exactly stages 0–3. With a bare Array a fifth depth entry renders the same bytes while
cnxDropTotal counts it, and a short one reads 0 through !.
deriving DecidableEq is what ConvNeXtRenderB's #guards on its restatement use, and
cnxModelName matches on the whole record rather than on the block count — which is what stops
B (36 blocks, like S) from introducing itself as an S.
Instances For
Equations
Instances For
Equations
Equations
- One or more equations did not get rendered due to their size.
Instances For
Equations
ConvNeXt-T: [3,3,9,3] @ [96,192,384,768], 28.6 M at K=1000. The DEFAULT everywhere, so
every pre-existing call site is unchanged and every committed artifact re-renders byte-identical.
Instances For
ConvNeXt-B: S's depth AND a wider net — [128,256,512,1024]. 88.6 M at K=1000.
B is the size that made the dims a parameter: it shares S's depth table exactly, so anything
keying on block count alone cannot tell them apart.
Instances For
Total ConvNeXt blocks = the reference's totalDrop, i.e. the ramp DENOMINATOR is this
minus 1. Derived from the stage table the render traverses, not restated beside it — which is
what makes it follow D to 36 at ConvNeXt-S without a second edit.
Equations
- Proofs.StableHLO.cnxDropTotal V = Vector.foldl (fun (x1 x2 : ℕ) => x1 + x2) 0 V.depths
Instances For
The ramp index of stage si's block j — the reference's dbi, which counts blocks over
the WHOLE net. The single source for the drop-site numbering: the forward calls it walking the
stages forwards, the backward walking them backwards, and they cannot disagree.
Equations
- Proofs.StableHLO.cnxBlockIdx si j V = List.foldl (fun (x1 x2 : ℕ) => x1 + x2) 0 (List.map (fun (k : ℕ) => V.depths[k]!) (List.range si)) + j
Instances For
The number of per-example drop-path scale inputs an SD ConvNeXt render takes: one per block.
Equations
Instances For
The model name a render's banner claims, DERIVED from the stage table it was rendered at.
A banner is a render's own description of what it did, and the one thing worse than not
having one is having one that lies — cnxDropFwdBanner's own note says exactly that about
reusing a drop-free banner on an SD artifact. Passing the name as a second parameter beside D
would be two writers for one fact. So it is read off D.
It matches on the WHOLE RECORD, not on the block count, and that is not fussiness.
B is [3,3,27,3] like S, so a name keyed on cnxDropTotal (36 blocks) names S and B alike and
every B artifact would open by calling itself a ConvNeXt-S. The guards below enumerate all
three rather than spot-checking one. An unrecognised table gets a name that SAYS it is
unrecognised, rather than falling back to "ConvNeXt-T" the way a getD would.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The %dp<i>: tensor<Bxf32> inputs an SD render appends to its signature — one per block, in
ramp-index order, which is the order the driver's dropScales writes them into the blob.
Empty when off, which is what keeps every committed artifact byte-identical.
Equations
Instances For
The channel LayerNorm. ConvNeXt's channel_layer_norm takes h·w statistics per
example, each over the c channels at one spatial position, with a per-channel [c] affine.
That is a different function from .bnF (bnForward n ε γ β: one mean and one variance over the
whole c·h·w map, scalar γ/β). Matching a reduction's literal across dimensions = [1] against
the reference's axis=1 does not check this: the artifact's tensor is rank-2 [B, c·h·w] and the
reference's is rank-4 NCHW.
No new SHlo op. ConvNeXt's channel-LN IS ViT's row-LN under a transpose — view an
example as [c, s] with s = h·w, transpose to [s, c], and each row is one spatial position
holding its c channels, which is exactly what rowLNFlat normalises. Checked on device
(lake build channel-ln): forward and all three backward pieces tie the closed form at rel 0, the
.bnF control fires at rel 0.82, and the transposes measure
free against the whole-net LN (XLA folds a transpose into the consumer's layout).
Nat multiplication is not definitionally associative, and the ambient index here is
c*h*h = (c*h)*h while the transpose needs c*(h*h). reassoc/unassoc (IndexCast) are
castIdx along Nat.mul_assoc; they are casts on the index, not on the value, so pretty walks the
same tree.
The %one/%zero constants the channel-LN chain binds lnRowF/lnRowBack's SCALAR γ/β to —
the real per-channel affine is rowScaleF/rowBiasF downstream, exactly as ViT does it.
Emitted once per module body.
Wire the operand and forget the prelude, and the artifact uses an SSA name nothing defines —
iree-compile/XLA say "use of undeclared SSA value name" and nothing before them says
anything. regen_verified_mlir.sh check's prelude audit is what catches it.
Equations
- One or more equations did not get rendered due to their size.
Instances For
One LayerNorm forward site — ConvNeXt's real channel-LN: transpose to [h·w, c],
normalise each spatial row over its channels at the scalar identities %one/%zero, apply
the real [c] affine (vecLnSite at m = h·w), transpose back.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The HEAD LN forward site — the paper's norm(x.mean([-2,-1])).
It is cnxLnFwdSite with the two transposes DELETED, and that is the whole difference: after
GAP the tensor is a single [d] row, so "normalise each spatial row over its channels" and
"normalise the feature vector" are the same function at m = 1. Mirrors
Proofs.StableHLO.headLNGraph op for op.
It is vecLnSite at m = 1, so every operand is annotated Vec (1 * d), never Vec d.
1 * d does NOT reduce for a VARIABLE d (Nat.mul recurses on its second argument), and d
here is V.dims[3]!, which moves with the ConvNeXt size. It is the same annotation
convNextBackAll's smoothing chain carries, for the same reason.
Equations
- Proofs.StableHLO.cnxHeadLnFwdSite cBS gN btN xin d = Proofs.StableHLO.vecLnSite cBS 1 d Proofs.StableHLO.cEPS✝ gN btN xin
Instances For
Equations
Every leaf below is one of the *Sgd/*Grad pairs whose dens differ by exactly θ − lr · ·
(*Sgd_eq_grad, all rfl) and whose emits differ by exactly the const-lr/multiply/subtract tail
(tests/TestBatchedEmitTie.lean, byte-PREFIX). So one traversal serves both renders.
ConvNeXt's param tails are factored out of the cotangent traversal — bwdBlock computes
cotangents and nothing else. lrStr is threaded but unused in adam mode: AdamW's learning rate
is the runtime %lr argument, not a baked literal.
Shapes are List Nat rather than rendered tensor<…> strings because the AdamW render needs
both forms: %{nm}/%{nm}m/%{nm}v for the moment slots and the raw dimensions for the emitted
Adam ops.
allParams is private (it is this file's internal signature source); this is the one thing
outside it that legitimately needs the list — tests/TestWdExcludeTie.lean, which must read
the SAME names and shapes the renderer chose %wd/%wdz from. Exposing an alias rather than
dropping private keeps the surface one definition wide.
Equations
- Proofs.StableHLO.cnxAllParams nClasses V = Proofs.StableHLO.allParams✝ nClasses V
Instances For
Every SSA name the ConvNeXt-T forward produces. convNextFwdText returns just logits;
the train steps additionally consume the block records, the two downsample names and the
head's gap/hn on the way back.
Instances For
Equations
@convnext_fwd from the per-example chain — convNextFwdChain, the forward both train
steps in this file emit, followed by return logits. Takes %x plus the 182 params of
allParams in func-arg order (183 inputs) and returns logits [cBS, nClasses]. Around the
chain the text is the signature, a banner comment and chLnPrelude.
The committed verified_mlir/convnext_fwd.mlir is written by ConvNeXtRenderB
(convNextFwdRenderB); TestConvNeXtFwdBTie checks this render against it byte for byte, and
scripts/regen_verified_mlir.sh check checks that its body is a prefix of
convnext_train_step.mlir's.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The whole-net forward + cotangent + backward traversal, SHARED by the SGD and AdamW renders.
Returns (code, params, softmax): every emitted line; one SSA per parameter — the updated
param at adam := false, the un-fused gradient at adam := true; and the softmax SSA,
which the AdamW render's report-only %loss reads.
smooth is the cotangent recipe. none → the SGD render's (softmax − onehot)/B, unchanged
down to the hand-written %dy divide. some (α, −α/K, B) → the label-smoothed
((softmax − onehot) + α·onehot − α/K)/B composed from kit ops, which is what the AdamW recipe
trains on. ConvNeXt already spells the ÷B explicitly (unlike ViT/R34, which fold the mean into
lr), so here the smoothing is the only difference between the two cotangents.
The softmax is prettyd on its own line rather than nested inside the .sub so that %loss
can read it; .operand is a leaf that emits nothing, so the fresh-name sequence is the same
either way.
Equations
- One or more equations did not get rendered due to their size.
Instances For
ConvNeXt-T (full [3,3,9,3]) SGD train step rendered from the verified AST, the writer of
verified_mlir/convnext_train_step.mlir: %x, the 182 params and %onehot in (184 inputs),
the 182 updated params out. All 182 parameter gradients are SHlo nodes, and every parameter
but the stem weight psW is updated by a fused SHlo SGD op whose output is the updated
parameter. Hand-written text: the signature, the %sc/%bsc constants, chLnPrelude, the
%dy divide, the GAP-backward block, and psW's update (sgdOf around the certified
.convStride4WeightGrad).
The cotangent is plain CE with an explicit ÷B — unlike ViT/R34, which fold the batch mean
into lr — so the committed cLR = 0.1 is an effective 0.1, the house convention spelled
differently.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The driver's variant slug for a given replica count: the artifact is
verified_mlir/convnext_<variant>_train_step.mlir, the entry point is
@convnext_<variant>_train_step, and LEAN_MLIR_VARIANT selects it. All three must agree —
the shim checks the entry name and refuses a mismatch outright ("entry mismatch") rather than
running the wrong graph. The #guards at the bottom pin the literal #eval paths against this.
Unlike mnv2AdamVariant/r34AdamVariant there is no batch suffix.
The EMA renders get their OWN slugs (ema/emadp), for the reason the RMSProp ones do: a
render carrying a fourth [θ|m|v|ema] region must never be able to overwrite the artifact the
AdamW trainer runs, whose blob has three: two writers for one path, where the last one wins,
and here it would also be an arity mismatch the driver could not survive.
Equations
- One or more equations did not get rendered due to their size.
Instances For
ConvNeXt AdamW train step rendered from the verified AST — the wrapper ConvNeXtRenderB
calls (with traversal set to its batched chain) to write convnext_adam_train_step.mlir and
the other ConvNeXt AdamW/EMA artifacts.
Same backward as convnext_train_step (convNextBackAll, one traversal) but taking the
un-fused gradients, each fed to the proven AdamW triple. The cotangent adds label smoothing
(α = 0.1, K = nClasses); ConvNeXt already spelled the ÷B explicitly, so that is the only
difference.
Interface at ConvNeXt-T, ema := false, no drop-path: 551 in (%x, 182 θ, 182 m, 182 v,
%lr/%bc1/%bc2, %onehot) / 549 out (182 θ', 182 m', 182 v', %loss/%bc1/%bc2), the
packed [θ|m|v] layout the AdamW trainer reads.
All 182 parameter gradients are SHlo nodes, and each update is the proven AdamW triple.
Hand-written text: the signature and constants, the GAP-backward block (in convNextBackAll),
the %bc1/%bc2 passthroughs, and the report-only %loss (on no gradient path).
replicas > 1 renders the DATA-PARALLEL variant to its own entry name and artifact path
via cnxAdamVariant, so producing it can never clobber the one the trainer runs. The only
difference is one all_reduce(add)/N per parameter gradient, between the certified gradient
and the certified AdamW triple: certified gradient → trusted collective → certified AdamW.
See optOne for the carve-out.
Equations
- One or more equations did not get rendered due to their size.