Module nn 

Minimal neural-network modules — exactly what GPT-2 needs.

Structs

Embedding

Token + positional embedding table. The token table is stored row-chunked so no single GPU binding exceeds the device limit (see roadmap pitfalls); on CPU it is a single chunk.

LayerNorm

Linear

y = x @ W + b, with W stored [in_features, out_features].