ViT-Style Model Helpers (API)

This module provides a compact, reusable ViT-style model constructor used by runnable examples.

This is intentionally minimal:

• patch embedding is a strided convolution,
• tokenization is a reshape + axis swap (N×C×H×W -> N×(H*W)×C),
• the “transformer” is a single encoder block,
• the head is a simple flatten + linear classifier.

The point is to keep examples readable while still exercising: Conv2d + tokenization + attention + FFN on both CPU and CUDA eager backends.

structure NN.API.nn.models.VitConfig : Type

Configuration for a small ViT-style classifier.

Shapes:

• input: N×C×H×W
• output: N×outDim

• batch : ℕ
• inC : ℕ
• inH : ℕ
• inW : ℕ
• patchH : ℕ
• patchW : ℕ
• stride : ℕ
• padding : ℕ
• dModel : ℕ
• outDim : ℕ
• numHeads : ℕ
• headDim : ℕ
• ffnHidden : ℕ

def NN.API.nn.models.instReprVitConfig.repr : VitConfig → ℕ → Std.Format

@[implicit_reducible]

instance NN.API.nn.models.instReprVitConfig : Repr VitConfig

def NN.API.nn.models.VitConfig.outH (cfg : VitConfig) : ℕ

def NN.API.nn.models.VitConfig.outW (cfg : VitConfig) : ℕ

def NN.API.nn.models.VitConfig.seqLen (cfg : VitConfig) : ℕ

def NN.API.nn.models.VitConfig.flatDim (cfg : VitConfig) : ℕ

@[reducible, inline]

abbrev NN.API.nn.models.vitInShape (cfg : VitConfig) : Shape

@[reducible, inline]

abbrev NN.API.nn.models.vitOutShape (cfg : VitConfig) : Shape

@[reducible, inline]

abbrev NN.API.nn.models.vitConvOutShape (cfg : VitConfig) : Shape

@[reducible, inline]

abbrev NN.API.nn.models.vitTokensShape (cfg : VitConfig) : Shape

def NN.API.nn.models.nchwToTokens (cfg : VitConfig) : LayerDef (vitConvOutShape cfg) (vitTokensShape cfg)

Patch-tokenization adapter: N×C×H×W -> N×(H*W)×C.

This is the “low-hanging fruit” to move out of examples: the reshape needs a small size proof.

def NN.API.nn.models.vit1 (cfg : VitConfig) (h_inC : cfg.inC ≠ 0 := by decide) (h_patchH : cfg.patchH ≠ 0 := by decide) (h_patchW : cfg.patchW ≠ 0 := by decide) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (vitInShape cfg) (vitOutShape cfg))

One-block ViT-style classifier.

This is the constructor used by torchlean vit. Keeping it here makes the example a one-liner: def mkModel := nn.models.vit1 cfg.