Indexed Models and Embeddings

This module defines models with non-differentiable tensor inputs, their scalar objectives, and embedding-table builders. Import NN.API.Neural.Indexed when constructing or executing an indexed model; ordinary sequential builders remain in NN.API.Neural.Builders.

structure TorchLean.nn.IndexedModel (σ τ : Shape) (β : Type) [Storage β] : Type 1

A shape-typed model with one non-differentiable tensor input.

• stateShapes : List Shape

  Shapes of trainable parameters and persistent buffers.

• kindValue : String
• initialStateValue : TorchLean.nn.State Float self.stateShapes
• initializationPlanValue : Option (Runtime.Autograd.Model.Module.RuntimeInit.Plan self.stateShapes)
• trainableMaskValue : Array Bool
• modelValidatorValue : Except String Unit
• inputValidatorValue : TorchLean.Tensor β σ → Except String Unit
• programValue (_mode : TorchLean.nn.Mode) {α : Type} [TorchLean.Storage α] [Context α] : Runtime.Autograd.Model.ProgramWithDataInputs α β self.stateShapes [σ] τ

opaque TorchLean.nn.IndexedModel.Internal.create {σ τ : Shape} {β : Type} [Storage β] (stateShapes : List Shape) (initialState : State Float stateShapes) (program : Mode → {α : Type} → [inst : Storage α] → [inst_1 : Context α] → Runtime.Autograd.Model.ProgramWithDataInputs α β stateShapes [σ] τ) (kind : String := "IndexedModel") (initializationPlan : Option (Module.RuntimeInit.Plan stateShapes) := none) (trainableMask : Array Bool := Array.replicate stateShapes.length true) (validateModel : Except String Unit := pure ()) (validateInput : Tensor β σ → Except String Unit := fun (x : Tensor β σ) => pure ()) : IndexedModel σ τ β

Construct an indexed model from its complete runtime definition.

def TorchLean.nn.IndexedModel.Internal.invalidConfiguration {β : Type} [Storage β] (input output : Shape) (kind message : String) : IndexedModel input output β

Build an indexed-model placeholder whose static configuration is known to be invalid.

The placeholder owns no parameters or buffers and is rejected by IndexedModel.validate before allocation or execution.

The readers below are opaque on purpose. IndexedModel keeps every field but stateShapes private, so the only way to look inside is through this namespace, and an opaque reader is one that simp, decide and rfl cannot unfold back into the field. Builders therefore stay free to change how a model is represented without any downstream proof noticing, which is the whole reason the fields were made private in the first place.

opaque TorchLean.nn.IndexedModel.Internal.kind {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : String

Label the builder gave this model, used in summaries and error messages.

opaque TorchLean.nn.IndexedModel.Internal.initialState {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : State Float model.stateShapes

Parameters and buffers the model starts from, in stateShapes order.

opaque TorchLean.nn.IndexedModel.Internal.initializationPlan {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : Option (Module.RuntimeInit.Plan model.stateShapes)

How the runtime should fill the state at allocation time, or none for an unplanned model.

opaque TorchLean.nn.IndexedModel.Internal.trainableMask {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : Array Bool

One flag per state entry saying whether the optimizer may update it.

opaque TorchLean.nn.IndexedModel.Internal.validateModel {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : Except String Unit

Static verdict on the configuration, checked before anything is allocated.

opaque TorchLean.nn.IndexedModel.Internal.validateInput {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) (input : Tensor β σ) : Except String Unit

Verdict on one concrete input, for conditions the shape type cannot express.

opaque TorchLean.nn.IndexedModel.Internal.program {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) (mode : Mode) {α : Type} [Storage α] [Context α] : Runtime.Autograd.Model.ProgramWithDataInputs α β model.stateShapes [σ] τ

The runtime program for a given Mode, scalar-polymorphic in its execution type.

def TorchLean.nn.IndexedModel.kind {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : String

Model label used in summaries and diagnostics.

def TorchLean.nn.IndexedModel.initialState {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : State Float model.stateShapes

Semantic initial values for the complete model state.

def TorchLean.nn.IndexedModel.requiresGrad {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : Array Bool

Gradient flags aligned with stateShapes.

def TorchLean.nn.IndexedModel.validate {σ τ : Shape} {β : Type} [Storage β] (model : IndexedModel σ τ β) : Except String Unit

Validate static model configuration before allocation or graph execution.

def TorchLean.nn.IndexedModel.andThen {β : Type} [Storage β] {σ τ υ : Shape} (first : IndexedModel σ τ β) (rest : Sequential τ υ) : IndexedModel σ υ β

Append an ordinary sequential model after an indexed-input model.

Scalar objectives

def TorchLean.nn.IndexedModel.Objective.fromLoss {β : Type} [Storage β] {σ τ υ : Shape} (model : IndexedModel σ τ β) (loss : {α : Type} → [inst : Storage α] → [inst_1 : Context α] → Runtime.Autograd.Model.Program α [τ, υ] []) (mode : Mode := Mode.train) : Module.ObjectiveDefinition β model.stateShapes [υ] [σ]

Pair an indexed-input model with a scalar loss.

The resulting training module accepts one ordinary target tensor followed by the model's non-differentiable input tensor. Keeping those packs separate ensures that indices cannot receive gradients or be reinterpreted through the model's floating-point element type.

The target shape is independent of the model output shape, so this constructor also supports losses whose labels use a different representation from the prediction. Training mode is the default.

def TorchLean.nn.IndexedModel.Objective.meanSquaredError {β : Type} [Storage β] {σ τ : Shape} (model : IndexedModel σ τ β) (reduction : Loss.Reduction := Loss.Reduction.mean) (mode : Mode := Mode.train) : Module.ObjectiveDefinition β model.stateShapes [τ] [σ]

Pair an indexed-input model with mean-squared error.

structure TorchLean.nn.Embedding (vocabularySize embeddingWidth : ℕ) : Type

A reusable trainable embedding table with vocabularySize rows and vectors of length embeddingWidth.

Unlike IndexedModel, this definition is independent of the eventual token-tensor shape. Calling table.model indices specializes it to an input shape with bounded Fin vocabularySize indices. Instantiating that model produces mutable weight storage; the definition itself remains immutable so it can be lowered and used in proofs.

• initialWeightValue : TorchLean.Tensor Float [vocabularySize, embeddingWidth]
• initializationPlanValue : Runtime.Autograd.Model.Module.RuntimeInit.Plan [[vocabularySize, embeddingWidth]]
• trainableValue : Bool
• validationValue : Except String Unit

opaque TorchLean.nn.Embedding.Internal.create {vocabularySize embeddingWidth : ℕ} (initialWeight : Tensor Float [vocabularySize, embeddingWidth]) (initializationPlan : Module.RuntimeInit.Plan [[vocabularySize, embeddingWidth]]) (trainable : Bool := true) (validation : Except String Unit := pure ()) : Embedding vocabularySize embeddingWidth

Construct an embedding table at the runtime-builder boundary.

opaque TorchLean.nn.Embedding.Internal.initialWeight {vocabularySize embeddingWidth : ℕ} (table : Embedding vocabularySize embeddingWidth) : Tensor Float [vocabularySize, embeddingWidth]

The table's starting weight matrix, one row per vocabulary entry.

opaque TorchLean.nn.Embedding.Internal.initializationPlan {vocabularySize embeddingWidth : ℕ} (table : Embedding vocabularySize embeddingWidth) : Module.RuntimeInit.Plan [[vocabularySize, embeddingWidth]]

How the runtime should fill that matrix at allocation time.

opaque TorchLean.nn.Embedding.Internal.isTrainable {vocabularySize embeddingWidth : ℕ} (table : Embedding vocabularySize embeddingWidth) : Bool

Whether the optimizer may update the table, false for a frozen lookup.

opaque TorchLean.nn.Embedding.Internal.validation {vocabularySize embeddingWidth : ℕ} (table : Embedding vocabularySize embeddingWidth) : Except String Unit

Static verdict on the table's configuration; Embedding.invalid is how it becomes an error.

def TorchLean.nn.Embedding.Internal.invalid (vocabularySize embeddingWidth : ℕ) (message : String) : Embedding vocabularySize embeddingWidth

Construct a rejected table value without exposing parameter state through Embedding.model.

def TorchLean.nn.Embedding.initialWeight {vocabularySize embeddingWidth : ℕ} (table : Embedding vocabularySize embeddingWidth) : Tensor Float [vocabularySize, embeddingWidth]

Initial table values, with one row per vocabulary item.

structure TorchLean.nn.Embedding.Config : Type

Construction options for a freshly initialized embedding table.

• weightInitialization : Init.Scheme

  Initialization scheme for the table.

  The default agrees with torch.nn.Embedding.reset_parameters: independent samples from the standard normal distribution. Language-model constructors normally override this with their architecture-specific initialization, such as GPT-2's standard deviation 0.02.

• freeze : Bool

  Freeze the table by excluding it from reverse-mode parameter gradients.

def TorchLean.nn.Embedding.Config.validateDimensions (vocabularySize embeddingWidth : ℕ) : Except String Unit

Validate dimensions shared by configured and exact-weight embeddings.

def TorchLean.nn.Embedding.Config.validate (config : Config) (vocabularySize embeddingWidth : ℕ) : Except String Unit

Validate embedding dimensions and initialization.

def TorchLean.nn.Embedding.fromWeight {vocabularySize embeddingWidth : ℕ} (weight : Tensor Float [vocabularySize, embeddingWidth]) (freeze : Bool := false) : Embedding vocabularySize embeddingWidth

Construct an embedding from an exact initial weight table.

This is the typed counterpart of passing _weight to torch.nn.Embedding. The supplied tensor fixes both dimensions at compile time and is also recorded as an exact row-major runtime initializer, so CPU and CUDA module construction start from the same payload.

def TorchLean.nn.Embedding.model {vocabularySize embeddingWidth : ℕ} (table : Embedding vocabularySize embeddingWidth) (input : Shape) : IndexedModel input (input.appendDim embeddingWidth) (Fin vocabularySize)

Specialize an embedding table to a concrete index-tensor shape.