Causal Transformer Runtime Support

Complete indexed-token and tied-weight causal Transformer models, together with their causal-language-model objective.

@[reducible, inline]

abbrev TorchLean.nn.models.CausalTransformer.Internal.tiedStateShapes (config : Config) {batchShape : Shape} (body : Sequential (config.embeddingShape batchShape) (config.embeddingShape batchShape)) : List Shape

State layout for a model whose token lookup and vocabulary projection share one matrix.

def TorchLean.nn.models.CausalTransformer.Internal.tiedForward {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (mode : Mode) (config : Config) {batchShape : Shape} (body : Sequential (config.embeddingShape batchShape) (config.embeddingShape batchShape)) (state : Runtime.Autograd.Torch.RefList (Runtime.ValueRef m α) (tiedStateShapes config body)) (tokens : Runtime.Autograd.Torch.DataRef (Fin config.vocabularySize) (config.tokenShape batchShape)) : m (Runtime.ValueRef m α (config.vocabularyShape batchShape))

Execute a tied token lookup, hidden Transformer, and vocabulary projection.

def TorchLean.nn.models.CausalTransformer.Internal.tiedModel (config : Config) {batchShape : Shape} (table : Embedding config.vocabularySize config.modelWidth) (body : Sequential (config.embeddingShape batchShape) (config.embeddingShape batchShape)) : IndexedModel (config.tokenShape batchShape) (config.vocabularyShape batchShape) (Fin config.vocabularySize)

Assemble a complete tied-weight indexed model from one table and one hidden body.

def TorchLean.nn.models.CausalTransformer.indexed (config : Config) (batchShape : Shape := []) : Builder (IndexedModel (config.tokenShape batchShape) (config.vocabularyShape batchShape) (Fin config.vocabularySize))

Build a complete indexed-token causal Transformer with an independent vocabulary head.

The token table, hidden Transformer, and output projection draw from one builder seed stream. The returned model is the single value used by training, checkpointing, and prediction.

def TorchLean.nn.models.CausalTransformer.tied (config : Config) (batchShape : Shape := []) : Builder (IndexedModel (config.tokenShape batchShape) (config.vocabularyShape batchShape) (Fin config.vocabularySize))

Build a complete causal Transformer with one matrix shared by token lookup and output projection.

def TorchLean.nn.models.CausalTransformer.objective (config : Config) {batchShape : Shape} (model : IndexedModel (config.tokenShape batchShape) (config.vocabularyShape batchShape) (Fin config.vocabularySize)) (reduction : Loss.Reduction := Loss.Reduction.mean) (mode : Mode := Mode.train) : Module.ObjectiveDefinition (Fin config.vocabularySize) model.stateShapes [] [config.tokenShape batchShape, config.tokenShape batchShape]

Pair a complete indexed-token causal Transformer with next-token cross entropy.

Inputs and targets are bounded token tensors with batchShape followed by the sequence axis.