ConvNeXt-T's step tie at the batched index, the un-fused gradient and the smoothed loss #

ConvNeXtStepTie.lean ties all 182 parameters of the SGD-inline convnext_train_step.mlir: each fused θ − lr·g op dens to the certified step at the cotangent the emitted backward chain delivers, per example, at a hard label. This file is that statement re-pointed along three axes; unlike EfficientNet-B0's (EfficientNetStepTieG.lean, two axes), ConvNeXt's per-example capstone was at a single image with the batch outside the AST, so the index is the third.

Axis 1 — the gradient node. Every conjunct is at the raw gradient node (*GradB), which is what convnext_adam_train_step.mlir and every f32 convnextin_* artifact emit; the fused op appears only in the SGD-inline file. The optimizer update that consumes the node (AdamW, the wx/clip variants, EMA) is outside this statement. ConvNeXtFoldGB.lean is the fold each conjunct delegates to.

Precision is a flag on the statement. The bf16 artifacts (convnextin_adamwxclipdropbf16 and its EMA, data-parallel, S and B twins) emit the bf16 conv, depthwise, strided-conv and patchify weight-gradient kinds (Foundation.Bf16GradNodes). Every such node below is stated on the renderers' switch — ConvWTiedBAt bf16 / DepthwiseWTiedBAt bf16 / ConvStridedWTiedBAt bf16 (Foundation.GradNodesBAt), the stem's convStride4WeightGradBAt bf16 id … spelled inline — bf16 := false the f32 node, bf16 := true the bf16 one, and the capstone cnx_net_tiedGB takes bf16, so it reaches the bf16 artifacts' gradient nodes read over ℝ exactly as it reads the f32 ones (Bf16Erasure: at the identity rounding the bf16 kind denotes what its f32 peer does). The bias nodes, LayerNorm, layer scale and the classifier carry no flag because no render switches them. The forward graph is a different matter: ConvNeXt's typed forward graph (convNextFwdGraphTCh) is the per-example one, tied to ConvNeXtRender, which has no bf16; the batched chain the bf16 artifacts render from has no typed graph at either precision.

Axis 2 — the loss. The capstone's top-of-chain cotangent is smoothedLossCotGraphDiv's, at a GENERAL target: the six-op chain expe → softmaxDiv → subB → scaleB → addVB → shiftB → divConstB this render emits, with the target arriving as the graph input %onehot — a soft vector under mixup or cutmix. The fused file pins it to softmax − oneHot, the gradient of plain cross-entropy at a hard label, which no ImageNet artifact computes. ConvNeXt's chain runs at the plain width N·K (no 1· row index), so there is no rowB/unrowB cast anywhere here.

Axis 3 — the index. N is a binder. Every forward activation is a per-example lift of the prefix the fused file threads (cnxStemFwdO, the block forward, cnxDownFwdChO), and every cotangent a per-example lift of the chain (the block's input cotangent, cnxDownCotInChAt, ConvNeXtChainClose's cnxCotP/E/N, chanLNTensor3Back); a block's lift is the INDEXED one (batchMapIdx / batchMapAuxIdx, Foundation.Batched.Indexed), since with stochastic depth example n runs the block at its own mask entry. That lift is honest for this net and for no BatchNorm net: LayerNorm, GELU, the convolutions, layer scale and the residual add are all batch-separable, so the batched op IS the per-example op under batchMap — which is what the *B constructors' den arms say. nC is a binder too: the Imagenette artifact is nC = 10, the ImageNet ones nC = 1000.

Relation to the fused file #

Every activation, every Jacobian witness and every chain cotangent is ConvNeXtStepTie.lean's, lifted; every conjunct's proof is one CnxFoldGB.*_den lemma. The wN, bN, gN, lrStr and lr binders disappear with the fused wrapper, as they did for B0. The head takes g as a parameter where the fused cnxHeadChTied computed it from a label; the per-block ties were already loss-agnostic.

Conventions carried unchanged from the fused file: channel LayerNorm (chanLNTensor3, a Vec c affine, h·w statistics per example) at all 22 spatial sites, the head at ViT's vector LN on one row, per-channel layer scale, GELU (no kink — no smoothness hypothesis anywhere), SYMMETRIC padding at the 4×4/s4 stem and the three 2×2/s2 downsamples. Stated at the literal widths of ConvNeXt-T; S and B are other nets.

Scope. One replica: in convnextin_adamdp* every gradient node feeds allReduceMeanF; DataParallel.Node composes the per-replica statement with the replica mean.

Stochastic depth is a binder. sd : Option (Fin 18 → Vec N) is the render's drop flag: none the drop-free artifacts' chain (every family constant, the statement the drop-free one by rfl), some the *drop* artifacts' — every shipped convnextin_*bf16 artifact, the book's convnextin_adamdpwxclipdroperfbf16 among them. Block k's site (cnxSd) sits between LayerScale and the skip add (CnxTieBlk.fwdOD), so every one of the block's nine nodes, LayerScale γ's included, reads the dropped cotangent (dropPathOpt, where ConvNeXtRenderB.bwdBlockB emits dropPathB on it) and only the skip fan-in the raw one (CnxTieBlk.cotInD, ConvNeXtDropBlock). The forward saves the nodes read are the drop-free ones: the site is after them. Note: the %dgi…%dgapf GAP backward is hand-written text on both chains; its value here is globalAvgPoolFlatHasVJP.backward, as in the fused file.

Per-example internal cotangents as functions of a block's INPUT #

ConvNeXtChainClose's cnxCotE/N and chanLNTensor3Back take the saved activations as arguments; batchMapAux lifts a function of (one saved value, one input), so these recompute the activations from the block input, exactly as cnxBlockCotInChAt does.

ConvNeXt block — all 9 gradient nodes, batched #

Downsample — channel-LN → 2×2/s2 conv, all 4 gradient nodes, batched #

Stem — 4×4/s4 patchify conv → channel-LN, all 4 gradient nodes, batched #

The bias grad is a pure cotangent reduce, so the render emits it as a stride-1 convBiasGradB at the OUTPUT resolution; the node carries an input it never reads (stated at zero), and the clause is at the real stride-4 stem on each example's image through GradNodeB.pdiv_flatConvStride4_bias_eq_conv2d. The weight is convStride4WeightGradB, at its gradient on both chains.

Head — GAP → vector-LN at one row → dense, all 4 gradient nodes, batched #

Stated at the LITERAL 768 for the fused file's reason: 1 * m does not reduce at a variable m. g is a PARAMETER — the loss cotangent arrives from smoothedLossCotGraphDiv in the capstone.

The stem wrapper — @[irreducible] #

Without it the capstone's refine ⟨cnx_stem_ch_tiedGBAt …, ?_, …⟩ times out in whnf at the default budget. The block, downsample and head ties need no wrapper: the capstone names cnxBlockChTiedGB / cnxDownChTiedGB / cnxHeadChTiedGB directly.

Every cotangent the capstone threads is a certified VJP backward #

The capstone below threads three per-example constructors: the head's cnxHeadDyXheadChN, each block's cotInD at its site and each downsample's cnxDownCotInChAt. Per example, each is the backward of its stage's certified VJP (CnxTieBlk.fwdODHasVJP, cnxDownChWHasVJP, and cnxHeadHasVJP below); lifted by batchMapAux_eq_batchMapHasVJPAt (the block's by batchMapAuxIdx_eq_batchMapIdxHasVJPAt), the batched cotangent is the backward of the lifted stage. ConvNeXt has no kink, so the only hypothesis is the LayerNorm's 0 < ε.

The whole-net capstone — all 182 params through the REAL batched forward + composed cotangent #

The fused file's thread, lifted: block inputs are per-example lifts of the forward prefixes (a block at its site, batchMapIdx), and the backward cotangents lifts of the per-example chain, composed from the smoothed loss g down through the head, every block's backward with the residual fan-in + dyOut at each of the eighteen identity-skip merges, the channel-LN-back at each of the three downsamples, and the stem LN's own back before the patchify conv's gradients.