Transformer

Transformer

Generated with Grok Build: Grok 4.5 · xAI Imagine (/imagine)

Modules under src/transformer/.

Public re-exports

pub use attention::MultiHeadAttention;
pub use block::TransformerBlock;
pub use model::{TransformerConfig, TransformerLM};

MultiHeadAttention

Scaled multi-head dot-product attention with causal masking. Weights are dense Tensors (Q/K/V/out projections). Constructed with new(dim, num_heads).

TransformerBlock

Pre-LayerNorm style residual stack: attention → residual → MLP (FeedForward) → residual. Constructed with new(dim, num_heads, mlp_dim).

TransformerLM / TransformerConfig

Decoder-only LM:

token embed + pos embed → N × TransformerBlock → final LN → LM head

API	Behavior
TransformerConfig::tiny()	Small config for tests
TransformerLM::new(cfg)	Random/init weights for forward smoke
forward(token_ids)	Logits [seq, vocab]
hidden_states(token_ids)	Post-block hidden before LM head (fusion hook)
param_count()	Estimated params

Requires seq_len ≤ max_seq_len.

Tests

test_tiny_forward, test_hidden_states in model.rs.

---

Last updated: July 29, 2026 Updated by: Grok Build: Grok 4.5 Package tip reference: 37c5a21 (main)