Documentation

LeanMlir.Proofs.Nets.ViT.ViTFwdGraph

ViT forward-graph pieces — the flat ↔ Mat bridges #

The ch10 token layer's ViT den helpers (StableHLO.Basic) are local re-spellings; this file ties them back to the proven Attention forms for the forward-graph faithfulness proofs (ViTMultiHead, ViTDepthK): each den helper applied to a Mat.flatten is the flatten of the Mat-level op, and patchEmbedF_x_den is the patch-embed stage every ViT forward-graph faithfulness proof starts from.

Flat ↔ Mat commutation bridges #

Each ch10 den helper applied to a Mat.flatten is the flatten of the corresponding Mat-level op (the Mat.unflatten_flatten round-trip cancels); the pointwise ops (scaleF/geluF/addV) commute with flattening definitionally. Public — ViTChainClose reuses them to tie the matmul-spelled SDPA backward to the proven closed forms.

theorem Proofs.StableHLO.rowLNFlat_flat {m n : ℕ} (ε γ β : ℝ) (A : Mat m n) :
rowLNFlat m n ε γ β A.flatten = Mat.flatten fun (r : Fin m) => layerNormForward n ε γ β (A r)
theorem Proofs.StableHLO.rowDenseFlat_flat {N a c : ℕ} (W : Mat a c) (b : Vec c) (A : Mat N a) :
rowDenseFlat N a c W b A.flatten = Mat.flatten fun (r : Fin N) => dense W b (A r)
theorem Proofs.StableHLO.matMulFlat_flat {m k n : ℕ} (A : Mat m k) (B : Mat k n) :
theorem Proofs.StableHLO.scale_flat {m n : ℕ} (s : ℝ) (A : Mat m n) :
(fun (i : Fin (m * n)) => s * A.flatten i) = Mat.flatten fun (r : Fin m) (c : Fin n) => s * A r c
theorem Proofs.StableHLO.gelu_flat {gf : GeluForm} {m n : ℕ} (A : Mat m n) :
gf.map (m * n) A.flatten = Mat.flatten fun (r : Fin m) => gf.map n (A r)
theorem Proofs.StableHLO.add_flat_pt {m n : ℕ} (A B : Mat m n) (j : Fin (m * n)) :
A.flatten j + B.flatten j = Mat.flatten (fun (r : Fin m) (s : Fin n) => A r s + B r s) j
theorem Proofs.StableHLO.patchEmbedF_x_den (ic H W patchSize N D : ℕ) (Wc : Kernel4 D ic patchSize patchSize) (bc cls : Vec D) (pos : Mat (N + 1) D) (x : Vec (ic * H * W)) :
den (SHlo.patchEmbedF "%Wp" "%bp" "%cls" "%pos" Wc bc cls pos (SHlo.operand "%x" x)) = (Mat.unflatten (patchEmbedFlat ic H W patchSize N D Wc bc cls pos x)).flatten

The patch-embed stage over the image operand %x, in the Mat.flatten (Mat.unflatten …) form the block den_aux lemmas take — stage 0 of every ViT forward-graph faithfulness proof.