Documentation

LeanMlir.Proofs.Codegen.StableHLO.Pretty

StableHLOPretty — the printer: SHlo terms → StableHLO text #

The syntactic half of StableHLO: pretty B g renders a term of the AST whose denotation den the proofs are about. SSA names are annotations den ignores, and skel erases the ℝ values, so the rendered text depends on the graph's shape and names only.

partwhat
ty / tyI1 / tyBf16 / tyF8, fresh, the ShapeTbl helperstype strings, SSA names, the name ↦ shape table
Raw, skelthe value-erased skeleton of an SHlo term
Tok, toToksthe post-order token stream
emitContract, emitTok, serializeToksStableHLO text per token (f32 / bf16 / fp8)
pretty, prettyAdamW, prettyAllReduceMean, renderModulewhole graphs and modules
fmt6 / fmt12, RmsHyper, bias-slot helpersthe literal and constant blocks the renderers share
*ModuleV, linTrainStepTextthe chapter 1–3 renderers

A module that states den facts imports StableHLO alone; one that renders text imports this. Trusted: the text emitTok prints for each token (operand order and arity, types, the lexical syntax) is not proved. scripts/gates/parse_verified_mlir.py parses every committed verified_mlir/*.mlir as StableHLO, and the trainers execute the artifacts they load under PJRT; StableHLOParse.roundtrip covers only the token encoding toToks, not the text.

Tensor-type string tensor<d₀x…xf32>.

Equations
Instances For

    Boolean (i1) tensor-type string, for compare/select masks.

    Equations
    Instances For

      bf16 tensor-type string: the convertF round node and the operands and results of the bf16 ops (dotInBf16, the …Bf16 convolutions).

      Equations
      Instances For

        fp8 peer of tyBf16. E4M3 only — f8E5M2 compiles, lowers to a plain __cublas$lt$matmul, and leaves ZERO f8e5m2 values in the optimized HLO: the type is silently widened away. Only E4M3 reaches the fp8 units on sm_89, so there is deliberately no E5M2 spelling here.

        Equations
        Instances For
          @[reducible, inline]

          SSA name ↦ the [c,h,w] the value bound to that name really carries. See liftPointwise.

          Keyed by NAME, not by flat width — and that is not a refinement, it is the whole correctness of the table. A width table collides whenever two layers have the same element count, and on the real nets they do: ConvNeXt-T's stage-2 MLP is 1536·14·14 = 301056 and its stage-0 block is 96·56·56 = 301056; stage 3's 3072·7·7 equals stage 1's 192·28·28. First writer won, so 24 of ConvNeXt's pointwise blocks unflattened to a shape with the right element count and the wrong layout — which is not a wrong program (the bracket is still an inverse reshape pair) but is exactly the relayout the bracket exists to remove. Keying by name cut ConvNeXt-T's transpose traffic by an order of magnitude and shortened its step (planning/archive/next_session_execution_and_parity.md).

          Newest entry first, and no dedup: fresh never reuses a name, so a lookup for a value the previous token produced hits the head of the list.

          The Bool is the value's LAYOUT: true means this is the map's row view [h·w, c] rather than the map [c, h, w]. It is not bookkeeping — it is what makes ConvNeXt's channel-LN transparent. That chain is transpose → lnRow → rowScale → rowBias → transpose, a layout ROUND TRIP whose two ends are the same [c,h,w] map; without the flag the closing transpose's result has no entry, the drop-path multiply that consumes it falls back to flat, and every pointwise op after it on the residual chain goes with it, which measurably grows ConvNeXt-T's relayout traffic (the same plan doc). And liftPointwise must NOT fire on a row view: [h·w, c] reshaped to [B,c,h,w] is a DIFFERENT permutation, not an inverse pair, so that one would be a wrong program rather than a slow one.

          Equations
          Instances For
            @[reducible, inline]

            Emitter state: the fresh-name counter, plus the name ↦ [c,h,w] table.

            The table lives in the STATE rather than in a pretty argument because a net renderer calls pretty once per graph FRAGMENT — a conv and the activation that consumes it land in different calls — and only the state is threaded across them.

            Equations
            Instances For

              Fresh SSA name %v{k}.

              Equations
              Instances For
                def Proofs.StableHLO.maxPool3s2FwdText (B c h w : ℕ) (r xn ninf p o : String) :

                The 3×3/s2 pool's emitted forward text, given already-freshened names.

                It is a shared helper rather than two copies for the reason sWGradGeom is: the per-example .maxPool3s2F and the batched BatchableOp.maxPool3s2 are two emitTok arms emitting one program, and a window or padding that drifted between them would be a pair of renders that agree on every structural check and compute different functions — which is the exact failure this whole op exists to fix. With one writer they cannot drift, and TestBatchedEmitTie then measures rather than assumes it.

                window_dimensions = 3, window_strides = 2, padding = [[1,1],[1,1]] on the spatial axes: He et al./torchvision MaxPool2d(3, stride=2, padding=1), window i = input [2i−1, 2i+1]. NOT XLA 'SAME', which pads (0,1) and slides the grid one input position — the two are different functions everywhere.

                Instances For
                  def Proofs.StableHLO.maxPool3s2BackText (B c h w : ℕ) (xN r xr dr z scn o : String) :

                  The 3×3/s2 pool's emitted backward text, given already-freshened names. Shared by the per-example and batched arms, for maxPool3s2FwdText's reason.

                  Only the window attributes differ from maxPoolBack's emit — nothing else — because select_and_scatter's scatter region already reduces with add, which is exactly the accumulation overlapping windows need.

                  %sa/%sb/%sc/%sd are hardcoded region block arguments and are therefore RESERVED SSA names: a top-level value of the same name is a redefinition error that surfaces only at XLA compile time.

                  Instances For

                    The stochastic-depth mask input name for ramp index i — the mName a dropPathB carries, and the tensor<Bxf32> the signature declares for it.

                    It lives HERE, beside the emitter, rather than in one net's renderer, because the spelling is load-bearing in three places that must agree and only one of them is Lean: dropPathP's emit reads it as an operand, every SD render's signature declares it, and scripts/probes/misplace_drop_sites.py matches %dp\d+ textually to build the placement control. A second definition would be the double-writer disease with a committed shell script as the third writer.

                    Equations
                    Instances For

                      The classifier-dropout mask input name — the mName a dropoutB carries, and the tensor<B×n×f32> the signature declares for it.

                      IT IS DELIBERATELY NOT %dp{i}-SHAPED, and that is not cosmetic. scripts/probes/misplace_drop_sites.py builds the stochastic-depth placement control by matching %dp\d+ textually; a dropout input spelled %dp9 would be swept into that rewrite, silently changing a control's meaning on a render it was never written for. A control that quietly does nothing reads exactly like a control that ran, and the cheap defence is a name the SD tooling cannot match. grep -c '%do' verified_mlir/*.mlir is 0 across every committed artifact.

                      Equations
                      Instances For

                        The renderable skeleton of an SHlo graph: opcodes + shapes + leaf SSA names, with ℝ operand values and the shape index erased — exactly what reaches the emitted text.

                        Instances For

                          Erase an SHlo graph to its renderable skeleton (drops ℝ values + shape index; keeps op structure, shapes, leaf names).

                          Instances For

                            One serialized token: an opcode with shapes/names; operands are positional.

                            Instances For
                              Equations
                              • One or more equations did not get rendered due to their size.
                              Instances For

                                Postorder serialization: children, then the node's opcode token.

                                Equations
                                Instances For

                                  Render a pointwise block at its 4-D shape when the OPERAND's producer recorded one. k receives the (possibly unflattened) input name and the dims to type its ops with, and returns (text, result name).

                                  The c*h*w == n guard is what keeps a mismatched entry from emitting an ill-typed reshape rather than merely a suboptimal one. It cannot fire today — an entry is written by the token that produced the name — and it is the difference between a missed optimisation and a render that does not parse, so it stays.

                                  The block's own RESULT is recorded too, which is what lets a pointwise CHAIN stay 4-D: the value crossing the token boundary keeps its flat type, so without this the second op in a swish→multiply→add chain would find nothing for its operand and drop back to flat.

                                  Equations
                                  • One or more equations did not get rendered due to their size.
                                  Instances For

                                    The text of the cross-replica mean — ViTRender.emitGradAllReduce's body, verbatim, so that the allReduceMean token re-renders every committed *dp* artifact byte-identically. all_reduce(add) over replica_groups = [[0..R-1]], then a divide by R; the names are %arsum{t} … %armean{t} from the tag rather than fresh. At R ≤ 1 there is no text and the operand's name is the result, exactly as the text function did.

                                    Equations
                                    • One or more equations did not get rendered due to their size.
                                    Instances For

                                      The squeeze-excite activations inside one seBlock emission, by name: the squeeze (GAP, [B,c]), the reduce dense + bias ([B,r]), its swish, and the excite dense + bias ([B,c]) — the four values the SE backward reads. k0 is the fresh-name counter when the seBlock token is emitted. The offsets are the order the "seBlock" arm of emitTok below calls fresh in (sq 5th, ex 8th, a1 10th, h2 13th); the #guard after emitTok checks each name is defined by the op it claims. A renderer that saves these reuses the SE the forward already computes instead of emitting a second, un-fused copy.

                                      Equations
                                      Instances For

                                        Render one token: pop its operands' result-names off the stack, emit its StableHLO line(s), push its fresh result name. The text each arm prints (operand order and arity, types, syntax) is trusted, checked by parsing and running the committed artifacts; the token stream it consumes is toToks (skel g), whose encoding round-trips (StableHLOParse.roundtrip).

                                        Instances For

                                          The bias's slot in a return-name list, gated the way biasName gates the operand: with convBias := false no bias SGD op is emitted, so the slot must LEAVE the list rather than carry the empty string the if convBias then … else pure ("", "") idiom hands back.

                                          This exists because leaving it in is silent twice over. An empty name renders return %a, , %b — malformed text, but only the lowerer ever sees it; and the name list keeps its FULL length, so an arity #guard on the signature still passes. Use this at every site where a names := [...] list is built from gated ops.

                                          Equations
                                          Instances For
                                            def Proofs.StableHLO.zeroBiasPrelude (convBias : Bool) (widths : List ℕ) :

                                            The zero-bias constants the convBias := false render consumes, one per channel width used as a conv bias. Emitted once at the top of the body; XLA folds the resulting add.

                                            Equations
                                            • One or more equations did not get rendered due to their size.
                                            Instances For

                                              Fixed-6-decimal float literal, so a computed smoothing constant emits in the SAME textual form the hand-written literals used and nClasses = 10 re-renders byte-identical.

                                              Equations
                                              • One or more equations did not get rendered due to their size.
                                              Instances For

                                                Fixed-12-decimal float literal, for constants fmt6 would destroy.

                                                It exists because fmt6 is not a formatting preference, it is a PRECISION CEILING, and small derived constants fall straight through it. Gradient accumulation's second-moment coefficient is (1−β₂)/K²; at K = 4 that is 6.25e-5, which fmt6 emits as 0.000063 — 0.8% wrong, in a baked literal, in the optimizer, where nothing downstream would question it. Same class as a hardcoded 0.010000 label-smoothing mass (see alphaOverK). fmt6 stays the default so every committed artifact re-renders byte-identically; this is for constants that need the room.

                                                Equations
                                                • One or more equations did not get rendered due to their size.
                                                Instances For
                                                  def Proofs.StableHLO.alphaOverK (nClasses : ℕ) (alpha : Float := 0.1) :

                                                  The label-smoothing mass per class, α/K. α = 0.1 throughout; K is nClasses.

                                                  Derived from nClasses, not written as a literal: a hardcoded 0.010000 is correct at K = 10 and wrong at every other K, and it sits in the COTANGENT, not just in the report-only %loss. At nClasses = 1000 it would make the smoothing term 100× too large — a different objective, silently. Nothing in the repo's proofs covers it: α is a literal in emitted text, so it needs its own numeric check.

                                                  Equations
                                                  Instances For

                                                    1 − α, the ON-class weight of label-smoothed CE. Emitted beside alphaOverK, because the two always move together and splitting them is how one of them gets updated alone.

                                                    Equations
                                                    Instances For

                                                      1 − ρ, the RMSProp mean-square mixing weight. Derived from ρ, never written as a second literal beside it — the oneMinusAlpha precedent, and the K-constant lesson: any emitted constant that depends on a hyperparameter must be DERIVED, because the copy is what gets left behind when the original moves.

                                                      Equations
                                                      Instances For

                                                        The RMSProp hyperparameters, as the JAX reference configs state them. ρ/μ are 0.9 on both nets that use this optimizer; ε and wd are what differ, and ε differs in the way that matters most (see Proofs.rmsBufNext_eps_placement_at_zero).

                                                        • rho : Float

                                                          rmspropDecay — the running mean-square decay.

                                                        • rho : Float

                                                          rmspropDecay — the running mean-square decay.

                                                        • mu : Float

                                                          momentum — μ for the buffer on the normalised gradient.

                                                        • mu : Float

                                                          momentum — μ for the buffer on the normalised gradient.

                                                        • eps : Float

                                                          rmspropEps — emitted INSIDE the square root (TensorFlow), not added to the root.

                                                        • eps : Float

                                                          rmspropEps — emitted INSIDE the square root (TensorFlow), not added to the root.

                                                        • wd : Float

                                                          COUPLED L2 (folded into the gradient), not AdamW's decoupled decay.

                                                        • wd : Float

                                                          COUPLED L2 (folded into the gradient), not AdamW's decoupled decay.

                                                        Instances For

                                                          ρ / (1−ρ) / μ / ε / wd as graph constants — the RMSProp peer of each renderer's adamConsts block. %lr stays a runtime tensor<f32> arg so one graph serves a whole LR schedule.

                                                          Equations
                                                          • One or more equations did not get rendered due to their size.
                                                          Instances For

                                                            MobileNetV2's RMSProp knobs (jax/MainMobilenetV2Imagenet.lean): ε = 1.0.

                                                            Equations
                                                            Instances For

                                                              EfficientNet-B0's RMSProp knobs (jax/MainEfficientNetImagenet.lean): ε = 1e-3.

                                                              Equations
                                                              Instances For

                                                                pretty — render an SHlo graph to StableHLO as serializeToks ∘ toToks ∘ skel: tokenize the graph (postorder), then print the tokens. The tokens it prints are exactly toToks (skel g), whose encoding StableHLOParse.roundtrip shows is invertible; how each token becomes text (emitTok) is trusted. A stack that does not end with one name returns "%MALFORMED" as the result name.

                                                                Equations
                                                                • One or more equations did not get rendered due to their size.
                                                                Instances For
                                                                  @[reducible, inline]

                                                                  The rounding a render hands the bf16/fp8 ops: the identity. A placeholder, exactly as a renderer's zero kernels are: a render produces TEXT, skel erases every ℝ payload before a token is emitted, and the emitted text is decided by the tag. The rounding-bearing den lives in the tie theorems, where rnd is arbitrary and the accuracy statement (Proofs/Float/*MixedFloatBridge.lean) instantiates it at bf16 round-to-nearest with |rnd x − x| ≤ 2⁻⁸|x|. A render that baked a concrete rounding here would be claiming the emitter knows about it, which it does not.

                                                                  Equations
                                                                  Instances For

                                                                    The cross-replica gradient mean as pretty of the allReduceMeanF node, used by every batched data-parallel render. At replicas ≤ 1 it emits nothing and threads the gradient's name. The R operand graphs are all .operand grad at a zero placeholder, because a render is value-independent — skel erases values — while the family is what den sums over in the tie.

                                                                    Equations
                                                                    • One or more equations did not get rendered due to their size.
                                                                    Instances For
                                                                      def Proofs.StableHLO.prettyAdamW (B : ℕ) (nm : String) (ds : List ℕ) (grad : String) (wdName : String := "%wd") :

                                                                      One parameter's AdamW update, as pretty of the proven triple — adamMNextF, adamVNextF and the decoupled-decay adamWParamF on the gradient grad, reading the graph constants adamWConsts binds, the step's %bc1/%bc2/%lr, and the decay operand wdName ("%wdz" for a timm no-decay parameter). Returns (code, θ', m', v'); every batched renderer's AdamW tail is this. wdName and den's wd must move together: both are 0 here only because every ℝ slot is a placeholder the emit ignores; if that changes, a no-decay site must pass wd := 0 as well as %wdz, or the artifact and the denotation describe different optimizers.

                                                                      Equations
                                                                      • One or more equations did not get rendered due to their size.
                                                                      Instances For
                                                                        def Proofs.StableHLO.adamWConsts (wdStr : String := "0.0001") :

                                                                        The AdamW graph constants — β₁ = 0.9, β₂ = 0.999, ε = 1e-8 and the baked decay wdStr (1e-4 is the Imagenette recipe every batched net shares; the ImageNet configs pass their own).

                                                                        Equations
                                                                        • One or more equations did not get rendered due to their size.
                                                                        Instances For
                                                                          def Proofs.StableHLO.renderModule (name argSig : String) (B retLen : ℕ) (g : SHlo retLen) :

                                                                          Wrap a rendered single-result graph as a func.func module.

                                                                          Equations
                                                                          • One or more equations did not get rendered due to their size.
                                                                          Instances For
                                                                            def Proofs.StableHLO.linTrainStepText (B m n : ℕ) (lrStr : String) (W : Mat m n) (b : Vec n) (x : Vec m) :

                                                                            The linear train step rendered ENTIRELY from the verified AST. Unlike linearTrainStepModuleV (forward via pretty, tail hand-written), here the whole module is pretty of denoted nodes: the cotangent (lossCotGraph, rendered once → shared %dy), then the two fused SGD ops weightSgd/biasSgd that consume %dy. So every emitted line is pretty(provenNode) and LinearFold proves the two outputs' den = the certified loss-descent SGD step. The lr ℝ / operand values are skel-erased (render is value-independent), so placeholders here render identically to the live graph the den theorems use.

                                                                            Equations
                                                                            • One or more equations did not get rendered due to their size.
                                                                            Instances For